Avoid Unnecessary Database Query on Back Relations (#275)

* Avoid Unnecessary Database Query on Back Relations

* Add `assign_related` option to `RelatedSet`; auto-assign back-relations for one-to-one

* Remove `not_nil!`

* Fix bool type

* Fix spec descriptions

spec/marten/db/field/many_to_one_spec.cr

@@ -212,5 +212,16 @@ describe Marten::DB::Field::ManyToOne do
 
       comment.article_id.should eq article.id!.hexstring
     end
+
+    it "retrieves the related article object for each comment" do
+      article = Marten::DB::Field::ManyToOneSpec::Article.create!(title: "This is an article", body: "This is a test")
+
+      Marten::DB::Field::ManyToOneSpec::Comment.create!(article: article, text: "This article is dope")
+      Marten::DB::Field::ManyToOneSpec::Comment.create!(article: article, text: "This article is not dope")
+
+      comments = Marten::DB::Field::ManyToOneSpec::Comment.all
+
+      comments[0].get_related_object_variable(:article)
+    end
   end
 end

spec/marten/db/field/one_to_one_spec.cr

@@ -280,5 +280,23 @@ describe Marten::DB::Field::OneToOne do
 
       comment.article_id.should eq article.id!.hexstring
     end
+
+    it "retrieves the related article object for the comment in a one-to-one association without bang operator" do
+      article = Marten::DB::Field::OneToOneSpec::Article.create!(title: "This is an article", body: "This is a test")
+      Marten::DB::Field::OneToOneSpec::Comment.create!(article: article, text: "This article is dope")
+
+      article = Marten::DB::Field::OneToOneSpec::Article.get(id: article.id).not_nil!
+
+      article.comment.not_nil!.get_related_object_variable(:article).should_not be_nil
+    end
+
+    it "retrieves the related article object for the comment in a one-to-one association with bang operator!" do
+      article = Marten::DB::Field::OneToOneSpec::Article.create!(title: "This is an article", body: "This is a test")
+      Marten::DB::Field::OneToOneSpec::Comment.create!(article: article, text: "This article is dope")
+
+      article = Marten::DB::Field::OneToOneSpec::Article.get!(id: article.id)
+
+      article.comment!.get_related_object_variable(:article).should_not be_nil
+    end
   end
 end

spec/marten/db/field/one_to_one_spec/models/comment.cr

@@ -1,7 +1,7 @@
 module Marten::DB::Field::OneToOneSpec
   class Comment < Marten::Model
     field :id, :uuid, primary_key: true
-    field :article, :one_to_one, to: Marten::DB::Field::OneToOneSpec::Article
+    field :article, :one_to_one, to: Marten::DB::Field::OneToOneSpec::Article, related: :comment
     field :text, :text
 
     after_initialize :initialize_id

spec/marten/db/query/related_set_spec.cr

@@ -99,4 +99,33 @@ describe Marten::DB::Query::RelatedSet do
       new_post.author.should eq user
     end
   end
+
+  describe "#fetch" do
+    it "assigns the related object if assign_related is explicitly set to true" do
+      user = TestUser.create!(username: "jd1", email: "jd1@example.com", first_name: "John", last_name: "Doe")
+      Post.create!(author: user, title: "Post 1")
+
+      qset = Marten::DB::Query::RelatedSet(Post).new(user, "author_id", assign_related: true)
+
+      qset[0].get_related_object_variable(:author).should eq user
+    end
+
+    it "does not assign the related object when assign_related is explicitly set to false" do
+      user = TestUser.create!(username: "jd1", email: "jd1@example.com", first_name: "John", last_name: "Doe")
+      Post.create!(author: user, title: "Post 1")
+
+      qset = Marten::DB::Query::RelatedSet(Post).new(user, "author_id", assign_related: false)
+
+      qset[0].get_related_object_variable(:author).should be_nil
+    end
+
+    it "does not assign the related object when assign_related is left as the default value (false)" do
+      user = TestUser.create!(username: "jd1", email: "jd1@example.com", first_name: "John", last_name: "Doe")
+      Post.create!(author: user, title: "Post 1")
+
+      qset = Marten::DB::Query::RelatedSet(Post).new(user, "author_id")
+
+      qset[0].get_related_object_variable(:author).should be_nil
+    end
+  end
 end

src/marten/db/field/many_to_one.cr

@@ -202,7 +202,7 @@ module Marten
 
                     def {{ related_field_name.id }}
                       @_reverse_m2o_{{ related_field_name.id }} ||=
-                        {{ model_klass }}::RelatedQuerySet.new(self, {{ field_id.stringify }})
+                        {{ model_klass }}::RelatedQuerySet.new(self, {{ field_id.stringify }}, assign_related: true)
                     end
                   end
                 end

src/marten/db/field/one_to_one.cr

@@ -210,10 +210,18 @@ module Marten
 
                     def {{ related_field_name.id }}
                       @_reverse_o2o_{{ related_field_name.id }} ||= {{ model_klass }}.get({{ field_id }}: pk)
+
+                      if reverse_o2o = @_reverse_o2o_{{ related_field_name.id }}
+                        reverse_o2o.{{ relation_attribute_name }} ||= self
+                      end
+
+                      @_reverse_o2o_{{ related_field_name.id }}
                     end
 
                     def {{ related_field_name.id }}!
                       @_reverse_o2o_{{ related_field_name.id }} ||= {{ model_klass }}.get!({{ field_id }}: pk)
+                      @_reverse_o2o_{{ related_field_name.id }}.not_nil!.{{ relation_attribute_name }} ||= self
+                      @_reverse_o2o_{{ related_field_name.id }}.not_nil!
                     end
                   end
                 end

src/marten/db/model/querying.cr

@@ -861,9 +861,10 @@ module Marten
                 def initialize(
                   @instance : Marten::DB::Model,
                   @related_field_id : ::String,
-                  query : Marten::DB::Query::SQL::Query({{ @type }})? = nil
+                  query : Marten::DB::Query::SQL::Query({{ @type }})? = nil,
+                  @assign_related : ::Bool = false
                 )
-                  super(@instance, @related_field_id, query)
+                  super(@instance, @related_field_id, query, @assign_related)
                 end
 
                 {% for queryset_id, block in MODEL_SCOPES[:custom] %}
@@ -876,7 +877,8 @@ module Marten
                   ::{{ @type }}::RelatedQuerySet.new(
                     instance: @instance,
                     related_field_id: @related_field_id,
-                    query: other_query.nil? ? @query.clone : other_query.not_nil!
+                    query: other_query.nil? ? @query.clone : other_query.not_nil!,
+                    assign_related: @assign_related
                   )
                 end
               end

src/marten/db/query/related_set.cr

@@ -5,7 +5,12 @@ module Marten
     module Query
       # Represents a query set resulting from a many-to-one relation.
       class RelatedSet(M) < Set(M)
-        def initialize(@instance : Marten::DB::Model, @related_field_id : String, query : SQL::Query(M)? = nil)
+        def initialize(
+          @instance : Marten::DB::Model,
+          @related_field_id : String,
+          query : SQL::Query(M)? = nil,
+          @assign_related : ::Bool = false
+        )
           @query = if query.nil?
                      q = SQL::Query(M).new
                      q.add_query_node(Node.new({@related_field_id => @instance.pk}))
@@ -15,6 +20,11 @@ module Marten
                    end
         end
 
+        protected def fetch
+          super
+          @result_cache.not_nil!.each { |r| r.assign_related_object(@instance, @related_field_id) } if @assign_related
+        end
+
         protected def build_record(**kwargs)
           record = M.new(**kwargs)
           record.assign_related_object(@instance, @related_field_id)
@@ -25,7 +35,8 @@ module Marten
           RelatedSet(M).new(
             instance: @instance,
             related_field_id: @related_field_id,
-            query: other_query.nil? ? @query.clone : other_query.not_nil!
+            query: other_query.nil? ? @query.clone : other_query.not_nil!,
+            assign_related: @assign_related
           )
         end
       end