From 49f6ccba4273f241e820e62b3391ae9d04774cd2 Mon Sep 17 00:00:00 2001
From: Jeremy Woertink <jeremywoertink@gmail.com>
Date: Fri, 29 Dec 2023 14:10:49 -0800
Subject: [PATCH] Adding block override for preloading with has_one. Fixes #985
 (#994)

---
 spec/avram/preloading/preloading_has_many_spec.cr | 11 +++++++++++
 spec/avram/preloading/preloading_has_one_spec.cr  | 12 ++++++++++++
 src/avram/associations/has_one.cr                 |  5 +++++
 3 files changed, 28 insertions(+)

diff --git a/spec/avram/preloading/preloading_has_many_spec.cr b/spec/avram/preloading/preloading_has_many_spec.cr
index 22f704c2b..739c8d5c2 100644
--- a/spec/avram/preloading/preloading_has_many_spec.cr
+++ b/spec/avram/preloading/preloading_has_many_spec.cr
@@ -51,6 +51,17 @@ describe "Preloading has_many associations" do
     end
   end
 
+  it "works with blocks" do
+    with_lazy_load(enabled: false) do
+      post = PostFactory.create
+      comment = CommentFactory.create &.post_id(post.id)
+
+      posts = Post::BaseQuery.new.preload_comments(&.id.not.eq(comment.id))
+
+      posts.results.first.comments.should eq([] of Comment)
+    end
+  end
+
   it "works with UUID foreign keys" do
     with_lazy_load(enabled: false) do
       item = LineItemFactory.create
diff --git a/spec/avram/preloading/preloading_has_one_spec.cr b/spec/avram/preloading/preloading_has_one_spec.cr
index 8ccacc869..019575998 100644
--- a/spec/avram/preloading/preloading_has_one_spec.cr
+++ b/spec/avram/preloading/preloading_has_one_spec.cr
@@ -34,6 +34,18 @@ describe "Preloading has_one associations" do
     end
   end
 
+  it "works with a block" do
+    with_lazy_load(enabled: false) do
+      admin = AdminFactory.create
+      sign_in_credential = SignInCredentialFactory.create &.user_id(admin.id)
+
+      admin = Admin::BaseQuery.new.preload_sign_in_credential(&.user_id(admin.id))
+
+      admin.first.sign_in_credential_preloaded?.should eq(true)
+      admin.first.sign_in_credential.should eq sign_in_credential
+    end
+  end
+
   it "works with optional association" do
     with_lazy_load(enabled: false) do
       UserFactory.create
diff --git a/src/avram/associations/has_one.cr b/src/avram/associations/has_one.cr
index ceac6f571..a859da55a 100644
--- a/src/avram/associations/has_one.cr
+++ b/src/avram/associations/has_one.cr
@@ -96,6 +96,11 @@ module Avram::Associations::HasOne
         preload_{{ assoc_name }}({{ model }}::BaseQuery.new)
       end
 
+      def preload_{{ assoc_name }} : self
+        modified_query = yield {{ model }}::BaseQuery.new
+        preload_{{ assoc_name }}(modified_query)
+      end
+
       def preload_{{ assoc_name }}(preload_query : {{ model }}::BaseQuery) : self
         add_preload do |records|
           ids = records.map(&.id)