activerecord 5.0.7.2 → 6.0.6.1

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

@@ -1,146 +1,150 @@
+# frozen_string_literal: true
+
 module ActiveRecord
   module Associations
     class Preloader
       class Association #:nodoc:
-        attr_reader :owners, :reflection, :preload_scope, :model, :klass
-        attr_reader :preloaded_records
-
-        def initialize(klass, owners, reflection, preload_scope)
+        def initialize(klass, owners, reflection, preload_scope, associate_by_default = true)
           @klass         = klass
-          @owners        = owners
+          @owners        = owners.uniq(&:__id__)
           @reflection    = reflection
           @preload_scope = preload_scope
+          @associate     = associate_by_default || !preload_scope || preload_scope.empty_scope?
           @model         = owners.first && owners.first.class
-          @scope         = nil
-          @preloaded_records = []
         end
 
-        def run(preloader)
-          preload(preloader)
-        end
+        def run
+          records = records_by_owner
 
-        def preload(preloader)
-          raise NotImplementedError
-        end
+          owners.each do |owner|
+            associate_records_to_owner(owner, records[owner] || [])
+          end if @associate
 
-        def scope
-          @scope ||= build_scope
+          self
         end
 
-        def records_for(ids)
-          query_scope(ids)
-        end
+        def records_by_owner
+          load_records unless defined?(@records_by_owner)
 
-        def query_scope(ids)
-          scope.where(association_key_name => ids)
+          @records_by_owner
         end
 
-        def table
-          klass.arel_table
-        end
+        def preloaded_records
+          load_records unless defined?(@preloaded_records)
 
-        # The name of the key on the associated records
-        def association_key_name
-          raise NotImplementedError
+          @preloaded_records
         end
 
-        # This is overridden by HABTM as the condition should be on the foreign_key column in
-        # the join table
-        def association_key
-          klass.arel_attribute(association_key_name, table)
-        end
+        private
+          attr_reader :owners, :reflection, :preload_scope, :model, :klass
 
-        # The name of the key on the model which declares the association
-        def owner_key_name
-          raise NotImplementedError
-        end
+          def load_records
+            # owners can be duplicated when a relation has a collection association join
+            # #compare_by_identity makes such owners different hash keys
+            @records_by_owner = {}.compare_by_identity
+            raw_records = owner_keys.empty? ? [] : records_for(owner_keys)
 
-        def options
-          reflection.options
-        end
+            @preloaded_records = raw_records.select do |record|
+              assignments = []
 
-        private
+              owners_by_key[convert_key(record[association_key_name])].each do |owner|
+                entries = (@records_by_owner[owner] ||= [])
+
+                if reflection.collection? || entries.empty?
+                  entries << record
+                  assignments << record
+                end
+              end
+
+              !assignments.empty?
+            end
+          end
 
-        def associated_records_by_owner(preloader)
-          records = load_records do |record|
-            owner = owners_by_key[convert_key(record[association_key_name])]
+          # The name of the key on the associated records
+          def association_key_name
+            reflection.join_primary_key(klass)
+          end
+
+          # The name of the key on the model which declares the association
+          def owner_key_name
+            reflection.join_foreign_key
+          end
+
+          def associate_records_to_owner(owner, records)
             association = owner.association(reflection.name)
-            association.set_inverse_instance(record)
+            if reflection.collection?
+              association.target = records
+            else
+              association.target = records.first
+            end
           end
 
-          owners.each_with_object({}) do |owner, result|
-            result[owner] = records[convert_key(owner[owner_key_name])] || []
+          def owner_keys
+            @owner_keys ||= owners_by_key.keys
           end
-        end
 
-        def owner_keys
-          unless defined?(@owner_keys)
-            @owner_keys = owners.map do |owner|
-              owner[owner_key_name]
+          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
-            @owner_keys.uniq!
-            @owner_keys.compact!
           end
-          @owner_keys
-        end
 
-        def owners_by_key
-          unless defined?(@owners_by_key)
-            @owners_by_key = owners.each_with_object({}) do |owner, h|
-              h[convert_key(owner[owner_key_name])] = owner
+          def key_conversion_required?
+            unless defined?(@key_conversion_required)
+              @key_conversion_required = (association_key_type != owner_key_type)
             end
+
+            @key_conversion_required
           end
-          @owners_by_key
-        end
 
-        def key_conversion_required?
-          @key_conversion_required ||= association_key_type != owner_key_type
-        end
+          def convert_key(key)
+            if key_conversion_required?
+              key.to_s
+            else
+              key
+            end
+          end
 
-        def convert_key(key)
-          if key_conversion_required?
-            key.to_s
-          else
-            key
+          def association_key_type
+            @klass.type_for_attribute(association_key_name).type
           end
-        end
 
-        def association_key_type
-          @klass.type_for_attribute(association_key_name.to_s).type
-        end
+          def owner_key_type
+            @model.type_for_attribute(owner_key_name).type
+          end
 
-        def owner_key_type
-          @model.type_for_attribute(owner_key_name.to_s).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 load_records(&block)
-          return {} if owner_keys.empty?
-          # Some databases impose a limit on the number of ids in a list (in Oracle it's 1000)
-          # Make several smaller queries if necessary or make one query if the adapter supports it
-          slices  = owner_keys.each_slice(klass.connection.in_clause_length || owner_keys.size)
-          @preloaded_records = slices.flat_map do |slice|
-            records_for(slice).load(&block)
+          def scope
+            @scope ||= build_scope
           end
-          @preloaded_records.group_by do |record| 
-            convert_key(record[association_key_name])
+
+          def reflection_scope
+            @reflection_scope ||= begin
+              reflection.join_scopes(klass.arel_table, klass.predicate_builder, klass).inject(&:merge!) || klass.unscoped
+            end
           end
-        end
 
-        def reflection_scope
-          @reflection_scope ||= reflection.scope_for(klass)
-        end
+          def build_scope
+            scope = klass.scope_for_association
 
-        def build_scope
-          scope = klass.scope_for_association
+            if reflection.type && !reflection.through_reflection?
+              scope.where!(reflection.type => model.polymorphic_name)
+            end
 
-          if reflection.type
-            scope.where!(reflection.type => model.base_class.sti_name)
+            scope.merge!(reflection_scope) unless reflection_scope.empty_scope?
+            scope.merge!(preload_scope) if preload_scope
+            scope
           end
-
-          scope.merge!(reflection_scope)
-          scope.merge!(preload_scope) if preload_scope != NULL_RELATION
-          scope
-        end
       end
     end
   end

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

@@ -1,114 +1,115 @@
+# frozen_string_literal: true
+
 module ActiveRecord
   module Associations
     class Preloader
-      module ThroughAssociation #:nodoc:
-        def through_reflection
-          reflection.through_reflection
+      class ThroughAssociation < Association # :nodoc:
+        PRELOADER = ActiveRecord::Associations::Preloader.new(associate_by_default: false)
+
+        def initialize(*)
+          super
+          @already_loaded = owners.first.association(through_reflection.name).loaded?
         end
 
-        def source_reflection
-          reflection.source_reflection
+        def preloaded_records
+          @preloaded_records ||= source_preloaders.flat_map(&:preloaded_records)
         end
 
-        def associated_records_by_owner(preloader)
-          preloader.preload(owners,
-                            through_reflection.name,
-                            through_scope)
+        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)
 
-          through_records = owners.map do |owner|
-            association = owner.association through_reflection.name
+          @records_by_owner = owners.each_with_object({}) do |owner, result|
+            through_records = through_records_by_owner[owner] || []
 
-            center = target_records_from_association(association)
-            [owner, Array(center)]
-          end
-
-          reset_association owners, through_reflection.name
+            if @already_loaded
+              if source_type = reflection.options[:source_type]
+                through_records = through_records.select do |record|
+                  record[reflection.foreign_type] == source_type
+                end
+              end
+            end
 
-          middle_records = through_records.flat_map { |(_,rec)| rec }
+            records = through_records.flat_map do |record|
+              source_records_by_owner[record]
+            end
 
-          preloaders = preloader.preload(middle_records,
-                                         source_reflection.name,
-                                         reflection_scope)
+            records.compact!
+            records.sort_by! { |rhs| preload_index[rhs] } if scope.order_values.any?
+            records.uniq! if scope.distinct_value
+            result[owner] = records
+          end
+        end
 
-          @preloaded_records = preloaders.flat_map(&:preloaded_records)
+        private
+          def source_preloaders
+            @source_preloaders ||= PRELOADER.preload(middle_records, source_reflection.name, scope)
+          end
 
-          middle_to_pl = preloaders.each_with_object({}) do |pl,h|
-            pl.owners.each { |middle|
-              h[middle] = pl
-            }
+          def middle_records
+            through_preloaders.flat_map(&:preloaded_records)
           end
 
-          through_records.each_with_object({}) do |(lhs,center), records_by_owner|
-            pl_to_middle = center.group_by { |record| middle_to_pl[record] }
+          def through_preloaders
+            @through_preloaders ||= PRELOADER.preload(owners, through_reflection.name, through_scope)
+          end
 
-            records_by_owner[lhs] = pl_to_middle.flat_map do |pl, middles|
-              rhs_records = middles.flat_map { |r|
-                association = r.association source_reflection.name
+          def through_reflection
+            reflection.through_reflection
+          end
 
-                target_records_from_association(association)
-              }.compact
+          def source_reflection
+            reflection.source_reflection
+          end
 
-              # Respect the order on `reflection_scope` if it exists, else use the natural order.
-              if reflection_scope.values[:order].present?
-                @id_map ||= id_to_index_map @preloaded_records
-                rhs_records.sort_by { |rhs| @id_map[rhs] }
-              else
-                rhs_records
-              end
+          def preload_index
+            @preload_index ||= preloaded_records.each_with_object({}).with_index do |(record, result), index|
+              result[record] = index
             end
           end
-        end
 
-        private
+          def through_scope
+            scope = through_reflection.klass.unscoped
+            options = reflection.options
 
-        def id_to_index_map(ids)
-          id_map = {}
-          ids.each_with_index { |id, index| id_map[id] = index }
-          id_map
-        end
+            values = reflection_scope.values
+            if annotations = values[:annotate]
+              scope.annotate!(*annotations)
+            end
 
-        def reset_association(owners, association_name)
-          should_reset = (through_scope != through_reflection.klass.unscoped) ||
-             (reflection.options[:source_type] && through_reflection.collection?)
+            if options[:source_type]
+              scope.where! reflection.foreign_type => options[:source_type]
+            elsif !reflection_scope.where_clause.empty?
+              scope.where_clause = reflection_scope.where_clause
 
-          # Don't cache the association - we would only be caching a subset
-          if should_reset
-            owners.each { |owner|
-              owner.association(association_name).reset
-            }
-          end
-        end
+              if includes = values[:includes]
+                scope.includes!(source_reflection.name => includes)
+              else
+                scope.includes!(source_reflection.name)
+              end
 
-        def through_scope
-          scope = through_reflection.klass.unscoped
-          values = reflection_scope.values
+              if values[:references] && !values[:references].empty?
+                scope.references!(values[:references])
+              else
+                scope.references!(source_reflection.table_name)
+              end
 
-          if options[:source_type]
-            scope.where! reflection.foreign_type => options[:source_type]
-          else
-            unless reflection_scope.where_clause.empty?
-              scope.includes_values = Array(values[:includes] || options[:source])
-              scope.where_clause = reflection_scope.where_clause
               if joins = values[:joins]
                 scope.joins!(source_reflection.name => joins)
               end
+
               if left_outer_joins = values[:left_outer_joins]
                 scope.left_outer_joins!(source_reflection.name => left_outer_joins)
               end
-            end
 
-            scope.references! values[:references]
-            if scope.eager_loading? && order_values = values[:order]
-              scope = scope.order(order_values)
+              if scope.eager_loading? && order_values = values[:order]
+                scope = scope.order(order_values)
+              end
             end
-          end
 
-          scope
-        end
-
-        def target_records_from_association(association)
-          association.loaded? ? association.target : association.reader
-        end
+            scope
+          end
       end
     end
   end