lib/torque/postgresql/associations/preloader/association.rb in torque-postgresql-2.0.3 vs lib/torque/postgresql/associations/preloader/association.rb in torque-postgresql-2.0.4

@@ -1,5 +1,7 @@
+# frozen_string_literal: true
+
 module Torque
   module PostgreSQL
     module Associations
       module Preloader
         module Association
@@ -36,14 +38,40 @@
               owners.each do |owner|
                 associate_records_to_owner(owner, records[owner[owner_key_name]] || [])
               end
             end
 
+            if PostgreSQL::AR610
+            # This is how Rails 6.1 now load the records
+              def load_records
+                return super unless connected_through_array?
+
+                @records_by_owner = {}.compare_by_identity
+                raw_records = owner_keys.empty? ? [] : records_for(owner_keys)
+
+                @preloaded_records = raw_records.select do |record|
+                  assignments = false
+
+                  ids = convert_key(record[association_key_name])
+                  owners_by_key.values_at(*ids).flat_map do |owner|
+                    entries = (@records_by_owner[owner] ||= [])
+
+                    if reflection.collection? || entries.empty?
+                      entries << record
+                      assignments = true
+                    end
+                  end
+
+                  assignments
+                end
+              end
+            end
+
             # Build correctly the constraint condition in order to get the
             # associated ids
             def records_for(ids, &block)
               return super unless connected_through_array?
-              condition = scope.arel_attribute(association_key_name)
+              condition = scope.arel_table[association_key_name]
               condition = reflection.build_id_constraint(condition, ids.flatten.uniq)
               scope.where(condition).load(&block)
             end
 
             def associate_records_to_owner(owner, records)