activerecord 6.0.0 → 7.2.3

data/lib/active_record/associations/preloader/association.rb

@@ -1,53 +1,241 @@
 # frozen_string_literal: true
 
+# :enddoc:
+
 module ActiveRecord
   module Associations
     class Preloader
-      class Association #:nodoc:
-        def initialize(klass, owners, reflection, preload_scope)
+      class Association # :nodoc:
+        class LoaderQuery
+          attr_reader :scope, :association_key_name
+
+          def initialize(scope, association_key_name)
+            @scope = scope
+            @association_key_name = association_key_name
+          end
+
+          def eql?(other)
+            association_key_name == other.association_key_name &&
+              scope.table_name == other.scope.table_name &&
+              scope.connection_specification_name == other.scope.connection_specification_name &&
+              scope.values_for_queries == other.scope.values_for_queries
+          end
+
+          def hash
+            [association_key_name, scope.table_name, scope.connection_specification_name, scope.values_for_queries].hash
+          end
+
+          def records_for(loaders)
+            LoaderRecords.new(loaders, self).records
+          end
+
+          def load_records_in_batch(loaders)
+            raw_records = records_for(loaders)
+
+            loaders.each do |loader|
+              loader.load_records(raw_records)
+              loader.run
+            end
+          end
+
+          def load_records_for_keys(keys, &block)
+            return [] if keys.empty?
+
+            if association_key_name.is_a?(Array)
+              query_constraints = Hash.new { |hsh, key| hsh[key] = Set.new }
+
+              keys.each_with_object(query_constraints) do |values_set, constraints|
+                association_key_name.zip(values_set).each do |key_name, value|
+                  constraints[key_name] << value
+                end
+              end
+
+              scope.where(query_constraints)
+            else
+              scope.where(association_key_name => keys)
+            end.load(&block)
+          end
+        end
+
+        class LoaderRecords
+          def initialize(loaders, loader_query)
+            @loader_query = loader_query
+            @loaders = loaders
+            @keys_to_load = Set.new
+            @already_loaded_records_by_key = {}
+
+            populate_keys_to_load_and_already_loaded_records
+          end
+
+          def records
+            load_records + already_loaded_records
+          end
+
+          private
+            attr_reader :loader_query, :loaders, :keys_to_load, :already_loaded_records_by_key
+
+            def populate_keys_to_load_and_already_loaded_records
+              loaders.each do |loader|
+                loader.owners_by_key.each do |key, owners|
+                  if loaded_owner = owners.find { |owner| loader.loaded?(owner) }
+                    already_loaded_records_by_key[key] = loader.target_for(loaded_owner)
+                  else
+                    keys_to_load << key
+                  end
+                end
+              end
+
+              @keys_to_load.subtract(already_loaded_records_by_key.keys)
+            end
+
+            def load_records
+              loader_query.load_records_for_keys(keys_to_load) do |record|
+                loaders.each { |l| l.set_inverse(record) }
+              end
+            end
+
+            def already_loaded_records
+              already_loaded_records_by_key.values.flatten
+            end
+        end
+
+        attr_reader :klass
+
+        def initialize(klass, owners, reflection, preload_scope, reflection_scope, associate_by_default)
           @klass         = klass
-          @owners        = owners
+          @owners        = owners.uniq(&:__id__)
           @reflection    = reflection
           @preload_scope = preload_scope
+          @reflection_scope = reflection_scope
+          @associate     = associate_by_default || !preload_scope || preload_scope.empty_scope?
           @model         = owners.first && owners.first.class
+          @run = false
         end
 
-        def run
-          if !preload_scope || preload_scope.empty_scope?
-            owners.each do |owner|
-              associate_records_to_owner(owner, records_by_owner[owner] || [])
-            end
+        def table_name
+          @klass.table_name
+        end
+
+        def future_classes
+          if run?
+            []
           else
-            # Custom preload scope is used and
-            # the association can not be marked as loaded
-            # Loading into a Hash instead
-            records_by_owner
+            [@klass]
           end
+        end
+
+        def runnable_loaders
+          [self]
+        end
+
+        def run?
+          @run
+        end
+
+        def run
+          return self if run?
+          @run = true
+
+          records = records_by_owner
+
+          owners.each do |owner|
+            associate_records_to_owner(owner, records[owner] || [])
+          end if @associate
+
           self
         end
 
         def records_by_owner
+          load_records unless defined?(@records_by_owner)
+
+          @records_by_owner
+        end
+
+        def preloaded_records
+          load_records unless defined?(@preloaded_records)
+
+          @preloaded_records
+        end
+
+        # The name of the key on the associated records
+        def association_key_name
+          reflection.join_primary_key(klass)
+        end
+
+        def loader_query
+          LoaderQuery.new(scope, association_key_name)
+        end
+
+        def owners_by_key
+          @owners_by_key ||= owners.each_with_object({}) do |owner, result|
+            key = derive_key(owner, owner_key_name)
+            (result[key] ||= []) << owner if key
+          end
+        end
+
+        def loaded?(owner)
+          owner.association(reflection.name).loaded?
+        end
+
+        def target_for(owner)
+          Array.wrap(owner.association(reflection.name).target)
+        end
+
+        def scope
+          @scope ||= build_scope
+        end
+
+        def set_inverse(record)
+          if owners = owners_by_key[derive_key(record, association_key_name)]
+            # Processing only the first owner
+            # because the record is modified but not an owner
+            association = owners.first.association(reflection.name)
+            association.set_inverse_instance(record)
+          end
+        end
+
+        def load_records(raw_records = nil)
           # owners can be duplicated when a relation has a collection association join
           # #compare_by_identity makes such owners different hash keys
-          @records_by_owner ||= preloaded_records.each_with_object({}.compare_by_identity) do |record, result|
-            owners_by_key[convert_key(record[association_key_name])].each do |owner|
-              (result[owner] ||= []) << record
+          @records_by_owner = {}.compare_by_identity
+          raw_records ||= loader_query.records_for([self])
+          @preloaded_records = raw_records.select do |record|
+            assignments = false
+
+            owners_by_key[derive_key(record, association_key_name)]&.each do |owner|
+              entries = (@records_by_owner[owner] ||= [])
+
+              if reflection.collection? || entries.empty?
+                entries << record
+                assignments = true
+              end
             end
+
+            assignments
           end
         end
 
-        def preloaded_records
-          return @preloaded_records if defined?(@preloaded_records)
-          @preloaded_records = owner_keys.empty? ? [] : records_for(owner_keys)
-        end
+        def associate_records_from_unscoped(unscoped_records)
+          return if unscoped_records.nil? || unscoped_records.empty?
+          return if !reflection_scope.empty_scope?
+          return if preload_scope && !preload_scope.empty_scope?
+          return if reflection.collection?
 
-        private
-          attr_reader :owners, :reflection, :preload_scope, :model, :klass
+          unscoped_records.select { |r| r[association_key_name].present? }.each do |record|
+            owners = owners_by_key[derive_key(record, association_key_name)]
+            owners&.each_with_index do |owner, i|
+              association = owner.association(reflection.name)
+              association.target = record
 
-          # The name of the key on the associated records
-          def association_key_name
-            reflection.join_primary_key(klass)
+              if i == 0 # Set inverse on first owner
+                association.set_inverse_instance(record)
+              end
+            end
           end
+        end
+
+        private
+          attr_reader :owners, :reflection, :preload_scope, :model
 
           # The name of the key on the model which declares the association
           def owner_key_name
@@ -55,25 +243,18 @@ module ActiveRecord
           end
 
           def associate_records_to_owner(owner, records)
+            return if loaded?(owner)
+
             association = owner.association(reflection.name)
+
             if reflection.collection?
-              association.target = records
+              not_persisted_records = association.target.reject(&:persisted?)
+              association.target = records + not_persisted_records
             else
               association.target = records.first
             end
           end
 
-          def owner_keys
-            @owner_keys ||= owners_by_key.keys
-          end
-
-          def owners_by_key
-            @owners_by_key ||= owners.each_with_object({}) do |owner, result|
-              key = convert_key(owner[owner_key_name])
-              (result[key] ||= []) << owner if key
-            end
-          end
-
           def key_conversion_required?
             unless defined?(@key_conversion_required)
               @key_conversion_required = (association_key_type != owner_key_type)
@@ -82,6 +263,14 @@ module ActiveRecord
             @key_conversion_required
           end
 
+          def derive_key(owner, key)
+            if key.is_a?(Array)
+              key.map { |k| convert_key(owner._read_attribute(k)) }
+            else
+              convert_key(owner._read_attribute(key))
+            end
+          end
+
           def convert_key(key)
             if key_conversion_required?
               key.to_s
@@ -98,22 +287,8 @@ module ActiveRecord
             @model.type_for_attribute(owner_key_name).type
           end
 
-          def records_for(ids)
-            scope.where(association_key_name => ids).load do |record|
-              # Processing only the first owner
-              # because the record is modified but not an owner
-              owner = owners_by_key[convert_key(record[association_key_name])].first
-              association = owner.association(reflection.name)
-              association.set_inverse_instance(record)
-            end
-          end
-
-          def scope
-            @scope ||= build_scope
-          end
-
           def reflection_scope
-            @reflection_scope ||= reflection.scope ? reflection.scope_for(klass.unscoped) : klass.unscoped
+            @reflection_scope ||= reflection.join_scopes(klass.arel_table, klass.predicate_builder, klass).inject(klass.unscoped, &:merge!)
           end
 
           def build_scope
@@ -123,9 +298,17 @@ module ActiveRecord
               scope.where!(reflection.type => model.polymorphic_name)
             end
 
-            scope.merge!(reflection_scope) if reflection.scope
-            scope.merge!(preload_scope) if preload_scope
-            scope
+            scope.merge!(reflection_scope) unless reflection_scope.empty_scope?
+
+            if preload_scope && !preload_scope.empty_scope?
+              scope.merge!(preload_scope)
+            end
+
+            cascade_strict_loading(scope)
+          end
+
+          def cascade_strict_loading(scope)
+            preload_scope&.strict_loading_value ? scope.strict_loading : scope
           end
       end
     end

data/lib/active_record/associations/preloader/batch.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module Associations
+    class Preloader
+      class Batch # :nodoc:
+        def initialize(preloaders, available_records:)
+          @preloaders = preloaders.reject(&:empty?)
+          @available_records = available_records.flatten.group_by { |r| r.class.base_class }
+        end
+
+        def call
+          branches = @preloaders.flat_map(&:branches)
+          until branches.empty?
+            loaders = branches.flat_map(&:runnable_loaders)
+
+            loaders.each { |loader| loader.associate_records_from_unscoped(@available_records[loader.klass.base_class]) }
+
+            if loaders.any?
+              future_tables = branches.flat_map do |branch|
+                branch.future_classes - branch.runnable_loaders.map(&:klass)
+              end.map(&:table_name).uniq
+
+              target_loaders = loaders.reject { |l| future_tables.include?(l.table_name)  }
+              target_loaders = loaders if target_loaders.empty?
+
+              group_and_load_similar(target_loaders)
+              target_loaders.each(&:run)
+            end
+
+            finished, in_progress = branches.partition(&:done?)
+
+            branches = in_progress + finished.flat_map(&:children)
+          end
+        end
+
+        private
+          attr_reader :loaders
+
+          def group_and_load_similar(loaders)
+            loaders.grep_v(ThroughAssociation).group_by(&:loader_query).each_pair do |query, similar_loaders|
+              query.load_records_in_batch(similar_loaders)
+            end
+          end
+      end
+    end
+  end
+end

data/lib/active_record/associations/preloader/branch.rb

@@ -0,0 +1,153 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module Associations
+    class Preloader
+      class Branch # :nodoc:
+        attr_reader :association, :children, :parent
+        attr_reader :scope, :associate_by_default
+        attr_writer :preloaded_records
+
+        def initialize(association:, children:, parent:, associate_by_default:, scope:)
+          @association = if association
+            begin
+              @association = association.to_sym
+            rescue NoMethodError
+              raise ArgumentError, "Association names must be Symbol or String, got: #{association.class.name}"
+            end
+          end
+          @parent = parent
+          @scope = scope
+          @associate_by_default = associate_by_default
+
+          @children = build_children(children)
+          @loaders = nil
+        end
+
+        def future_classes
+          (immediate_future_classes + children.flat_map(&:future_classes)).uniq
+        end
+
+        def immediate_future_classes
+          if parent.done?
+            loaders.flat_map(&:future_classes).uniq
+          else
+            likely_reflections.reject(&:polymorphic?).flat_map do |reflection|
+              reflection.
+                chain.
+                map(&:klass)
+            end.uniq
+          end
+        end
+
+        def target_classes
+          if done?
+            preloaded_records.map(&:klass).uniq
+          elsif parent.done?
+            loaders.map(&:klass).uniq
+          else
+            likely_reflections.reject(&:polymorphic?).map(&:klass).uniq
+          end
+        end
+
+        def likely_reflections
+          parent_classes = parent.target_classes
+          parent_classes.filter_map do |parent_klass|
+            parent_klass._reflect_on_association(@association)
+          end
+        end
+
+        def root?
+          parent.nil?
+        end
+
+        def source_records
+          @parent.preloaded_records
+        end
+
+        def preloaded_records
+          @preloaded_records ||= loaders.flat_map(&:preloaded_records)
+        end
+
+        def done?
+          root? || (@loaders && @loaders.all?(&:run?))
+        end
+
+        def runnable_loaders
+          loaders.flat_map(&:runnable_loaders).reject(&:run?)
+        end
+
+        def grouped_records
+          h = {}
+          polymorphic_parent = !root? && parent.polymorphic?
+          source_records.each do |record|
+            reflection = record.class._reflect_on_association(association)
+            next if polymorphic_parent && !reflection || !record.association(association).klass
+            (h[reflection] ||= []) << record
+          end
+          h
+        end
+
+        def preloaders_for_reflection(reflection, reflection_records)
+          reflection_records.group_by do |record|
+            klass = record.association(association).klass
+
+            if reflection.scope && reflection.scope.arity != 0
+              # For instance dependent scopes, the scope is potentially
+              # different for each record. To allow this we'll group each
+              # object separately into its own preloader
+              reflection_scope = reflection.join_scopes(klass.arel_table, klass.predicate_builder, klass, record).inject(&:merge!)
+            end
+
+            [klass, reflection_scope]
+          end.map do |(rhs_klass, reflection_scope), rs|
+            preloader_for(reflection).new(rhs_klass, rs, reflection, scope, reflection_scope, associate_by_default)
+          end
+        end
+
+        def polymorphic?
+          return false if root?
+          return @polymorphic if defined?(@polymorphic)
+
+          @polymorphic = source_records.any? do |record|
+            reflection = record.class._reflect_on_association(association)
+            reflection && reflection.options[:polymorphic]
+          end
+        end
+
+        def loaders
+          @loaders ||=
+            grouped_records.flat_map do |reflection, reflection_records|
+              preloaders_for_reflection(reflection, reflection_records)
+            end
+        end
+
+        private
+          def build_children(children)
+            Array.wrap(children).flat_map { |association|
+              Array(association).flat_map { |parent, child|
+                Branch.new(
+                  parent: self,
+                  association: parent,
+                  children: child,
+                  associate_by_default: associate_by_default,
+                  scope: scope
+                )
+              }
+            }
+          end
+
+          # Returns a class containing the logic needed to load preload the data
+          # and attach it to a relation. The class returned implements a `run` method
+          # that accepts a preloader.
+          def preloader_for(reflection)
+            if reflection.options[:through]
+              ThroughAssociation
+            else
+              Association
+            end
+          end
+      end
+    end
+  end
+end

data/lib/active_record/associations/preloader/through_association.rb

@@ -4,26 +4,20 @@ module ActiveRecord
   module Associations
     class Preloader
       class ThroughAssociation < Association # :nodoc:
-        PRELOADER = ActiveRecord::Associations::Preloader.new
-
-        def initialize(*)
-          super
-          @already_loaded = owners.first.association(through_reflection.name).loaded?
-        end
-
         def preloaded_records
           @preloaded_records ||= source_preloaders.flat_map(&:preloaded_records)
         end
 
         def records_by_owner
-          return @records_by_owner if defined?(@records_by_owner)
-          source_records_by_owner = source_preloaders.map(&:records_by_owner).reduce(:merge)
-          through_records_by_owner = through_preloaders.map(&:records_by_owner).reduce(:merge)
+          @records_by_owner ||= owners.each_with_object({}) do |owner, result|
+            if loaded?(owner)
+              result[owner] = target_for(owner)
+              next
+            end
 
-          @records_by_owner = owners.each_with_object({}) do |owner, result|
             through_records = through_records_by_owner[owner] || []
 
-            if @already_loaded
+            if owners.first.association(through_reflection.name).loaded?
               if source_type = reflection.options[:source_type]
                 through_records = through_records.select do |record|
                   record[reflection.foreign_type] == source_type
@@ -42,17 +36,47 @@ module ActiveRecord
           end
         end
 
+        def runnable_loaders
+          if data_available?
+            [self]
+          elsif through_preloaders.all?(&:run?)
+            source_preloaders.flat_map(&:runnable_loaders)
+          else
+            through_preloaders.flat_map(&:runnable_loaders)
+          end
+        end
+
+        def future_classes
+          if run?
+            []
+          elsif through_preloaders.all?(&:run?)
+            source_preloaders.flat_map(&:future_classes).uniq
+          else
+            through_classes = through_preloaders.flat_map(&:future_classes)
+            source_classes = source_reflection.
+              chain.
+              reject { |reflection| reflection.respond_to?(:polymorphic?) && reflection.polymorphic? }.
+              map(&:klass)
+            (through_classes + source_classes).uniq
+          end
+        end
+
         private
+          def data_available?
+            owners.all? { |owner| loaded?(owner) } ||
+              through_preloaders.all?(&:run?) && source_preloaders.all?(&:run?)
+          end
+
           def source_preloaders
-            @source_preloaders ||= PRELOADER.preload(middle_records, source_reflection.name, scope)
+            @source_preloaders ||= ActiveRecord::Associations::Preloader.new(records: middle_records, associations: source_reflection.name, scope: scope, associate_by_default: false).loaders
           end
 
           def middle_records
-            through_preloaders.flat_map(&:preloaded_records)
+            through_records_by_owner.values.flatten
           end
 
           def through_preloaders
-            @through_preloaders ||= PRELOADER.preload(owners, through_reflection.name, through_scope)
+            @through_preloaders ||= ActiveRecord::Associations::Preloader.new(records: owners, associations: through_reflection.name, scope: through_scope, associate_by_default: false).loaders
           end
 
           def through_reflection
@@ -63,6 +87,14 @@ module ActiveRecord
             reflection.source_reflection
           end
 
+          def source_records_by_owner
+            @source_records_by_owner ||= source_preloaders.map(&:records_by_owner).reduce(:merge)
+          end
+
+          def through_records_by_owner
+            @through_records_by_owner ||= through_preloaders.map(&:records_by_owner).reduce(:merge)
+          end
+
           def preload_index
             @preload_index ||= preloaded_records.each_with_object({}).with_index do |(record, result), index|
               result[record] = index
@@ -73,6 +105,8 @@ module ActiveRecord
             scope = through_reflection.klass.unscoped
             options = reflection.options
 
+            return scope if options[:disable_joins]
+
             values = reflection_scope.values
             if annotations = values[:annotate]
               scope.annotate!(*annotations)
@@ -90,7 +124,7 @@ module ActiveRecord
               end
 
               if values[:references] && !values[:references].empty?
-                scope.references!(values[:references])
+                scope.references_values |= values[:references]
               else
                 scope.references!(source_reflection.table_name)
               end
@@ -108,7 +142,7 @@ module ActiveRecord
               end
             end
 
-            scope
+            cascade_strict_loading(scope)
           end
       end
     end