feat: reduce redundant queries on cross-db belongs_to (#734)

Author: matthv

app/services/forest_liana/base_getter.rb

@@ -21,11 +21,11 @@ def compute_includes
       @includes = ForestLiana::QueryHelper.get_one_association_names_symbol(@resource)
     end
 
-    def optimize_record_loading(resource, records)
+    def optimize_record_loading(resource, records, force_preload = true)
       polymorphic, preload_loads = analyze_associations(resource)
       result = records.eager_load(@includes.uniq - preload_loads - polymorphic)
 
-      result = result.preload(preload_loads) if Rails::VERSION::MAJOR >= 7
+      result = result.preload(preload_loads) if Rails::VERSION::MAJOR >= 7 && force_preload
 
       result
     end

app/services/forest_liana/resources_getter.rb

@@ -37,10 +37,10 @@ def perform
       end
 
       if includes.empty? || has_smart_fields
-        @records = optimize_record_loading(@resource, @records)
+        @records = optimize_record_loading(@resource, @records, false)
       else
         select = compute_select_fields
-        @records = optimize_record_loading(@resource, @records).references(includes).select(*select)
+        @records = optimize_record_loading(@resource, @records, false).references(includes).select(*select)
       end
 
       @records
@@ -55,7 +55,79 @@ def query_for_batch
     end
 
     def records
-      @records.offset(offset).limit(limit).to_a
+      records = @records.offset(offset).limit(limit).to_a
+
+      polymorphic_association, preload_loads = analyze_associations(@resource)
+      if polymorphic_association && Rails::VERSION::MAJOR >= 7
+        # TODO
+      end
+
+      preload_cross_database_associations(records, preload_loads)
+
+      records
+    end
+
+    def preload_cross_database_associations(records, preload_loads)
+      preload_loads.each do |association_name|
+        association = @resource.reflect_on_association(association_name)
+        next unless separate_database?(@resource, association)
+
+        columns = columns_for_cross_database_association(association_name)
+
+        if association.macro == :belongs_to
+          foreign_key = association.foreign_key
+          primary_key = association.klass.primary_key
+
+          ids = records.map { |r| r.public_send(foreign_key) }.compact.uniq
+          next if ids.empty?
+
+          associated = association.klass.where(primary_key => ids)
+                                  .select(columns)
+                                  .index_by { |record| record.public_send(primary_key) }
+
+          records.each do |record|
+            record.define_singleton_method(association_name) do
+              associated[record.send(foreign_key.to_sym)] || nil
+            end
+          end
+        end
+
+        if association.macro == :has_one
+          foreign_key = association.foreign_key
+          primary_key = association.active_record_primary_key
+
+          ids = records.map { |r| r.public_send(primary_key) }.compact.uniq
+          next if ids.empty?
+
+          associated = association.klass.where(foreign_key => ids)
+                                  .select(columns)
+                                  .index_by { |record| record.public_send(foreign_key.to_sym) }
+
+          records.each do |record|
+            record.define_singleton_method(association_name) do
+              associated[record.send(primary_key.to_sym)] || nil
+            end
+          end
+        end
+      end
+    end
+
+    def columns_for_cross_database_association(association_name)
+      return [:id] unless @params[:fields].present?
+
+      fields = @params[:fields][association_name.to_s]
+      return [:id] unless fields
+
+      base_fields = fields.split(',').map(&:strip).map(&:to_sym) | [:id]
+
+      association = @resource.reflect_on_association(association_name)
+      extra_key = association.foreign_key
+
+      # Add the foreign key used for the association to ensure it's available in the preloaded records
+      # This is necessary for has_one associations, without it calling record.public_send(foreign_key) would raise a "missing attribute" error
+      base_fields << extra_key if association.macro == :has_one
+
+      base_fields.uniq
     end
 
     def compute_includes