Refactor query constraints feature set

activerecord/lib/active_record/associations/builder/association.rb

@@ -19,7 +19,7 @@ class << self
     self.extensions = []
 
     VALID_OPTIONS = [
-      :class_name, :anonymous_class, :primary_key, :foreign_key, :dependent, :validate, :inverse_of, :strict_loading
+      :class_name, :anonymous_class, :primary_key, :foreign_key, :dependent, :validate, :inverse_of, :strict_loading, :query_constraints
     ].freeze # :nodoc:
 
     def self.build(model, name, scope, options, &block)

activerecord/lib/active_record/autosave_association.rb

@@ -505,7 +505,7 @@ def save_belongs_to_association(reflection)
             saved = record.save(validate: !autosave) if record.new_record? || (autosave && record.changed_for_autosave?)
 
             if association.updated?
-              primary_key = Array(reflection.options[:primary_key] || record.class.query_constraints_list)
+              primary_key = Array(compute_primary_key(reflection, record))
               foreign_key = Array(reflection.foreign_key)
 
               primary_key_foreign_key_pairs = primary_key.zip(foreign_key)
@@ -521,6 +521,16 @@ def save_belongs_to_association(reflection)
         end
       end
 
+      def compute_primary_key(reflection, record)
+        if primary_key_options = reflection.options[:primary_key]
+          primary_key_options
+        elsif query_constraints = record.class.query_constraints_list
+          query_constraints
+        else
+          :id
+        end
+      end
+
       def custom_validation_context?
         validation_context && [:create, :update].exclude?(validation_context)
       end

activerecord/lib/active_record/persistence.rb

@@ -492,7 +492,11 @@ def query_constraints(*columns_list)
       end
 
       def query_constraints_list # :nodoc:
-        @_query_constraints_list ||= query_constraints_list_fallback
+        @query_constraints_list ||= if base_class? || primary_key != base_class.primary_key
+          @_query_constraints_list
+        else
+          base_class.query_constraints_list
+        end
       end
 
       # Destroy an object (or multiple objects) that has the given id. The object is instantiated first,
@@ -632,17 +636,6 @@ def build_default_constraint
           default_where_clause = default_scoped(all_queries: true).where_clause
           default_where_clause.ast unless default_where_clause.empty?
         end
-
-        # This is a fallback method that is used to determine the query_constraints_list
-        # for cases when the model is not explicitly configured with query_constraints.
-        # For a base class, just use the primary key.
-        # For a child class, use the primary key unless primary key was overridden.
-        # If the child's primary key was not overridden, use the parent's query_constraints_list.
-        def query_constraints_list_fallback # :nodoc:
-          return Array(primary_key) if base_class? || primary_key != base_class.primary_key
-
-          base_class.query_constraints_list
-        end
     end
 
     # Returns true if this object hasn't been saved yet -- that is, a record
@@ -1144,8 +1137,12 @@ def _find_record(options)
     end
 
     def _in_memory_query_constraints_hash
-      self.class.query_constraints_list.index_with do |column_name|
-        attribute(column_name)
+      if self.class.query_constraints_list.nil?
+        { @primary_key => id }
+      else
+        self.class.query_constraints_list.index_with do |column_name|
+          attribute(column_name)
+        end
       end
     end
 
@@ -1155,8 +1152,12 @@ def apply_scoping?(options)
     end
 
     def _query_constraints_hash
-      self.class.query_constraints_list.index_with do |column_name|
-        attribute_in_database(column_name)
+      if self.class.query_constraints_list.nil?
+        { @primary_key => id_in_database }
+      else
+        self.class.query_constraints_list.index_with do |column_name|
+          attribute_in_database(column_name)
+        end
       end
     end
 

activerecord/lib/active_record/reflection.rb

@@ -490,9 +490,9 @@ def join_table
       end
 
       def foreign_key
-        @foreign_key ||= if options[:foreign_key] && options[:foreign_key].is_a?(Array)
+        @foreign_key ||= if options[:query_constraints]
           # composite foreign keys support
-          options[:foreign_key].map { |fk| fk.to_s.freeze }.freeze
+          options[:query_constraints].map { |fk| fk.to_s.freeze }.freeze
         else
           -(options[:foreign_key]&.to_s || derive_foreign_key)
         end
@@ -507,7 +507,7 @@ def association_primary_key(klass = nil)
       end
 
       def active_record_primary_key
-        @active_record_primary_key ||= if options[:foreign_key] && options[:foreign_key].is_a?(Array)
+        @active_record_primary_key ||= if options[:query_constraints]
           active_record.query_constraints_list
         else
           -(options[:primary_key]&.to_s || primary_key(active_record))
@@ -775,7 +775,7 @@ def association_class
 
       # klass option is necessary to support loading polymorphic associations
       def association_primary_key(klass = nil)
-        if options[:foreign_key] && options[:foreign_key].is_a?(Array)
+        if options[:query_constraints]
           (klass || self.klass).query_constraints_list
         elsif primary_key = options[:primary_key]
           @association_primary_key ||= -primary_key.to_s

activerecord/lib/active_record/relation/finder_methods.rb

@@ -576,15 +576,24 @@ def find_last(limit)
       end
 
       def ordered_relation
-        if order_values.empty? && (implicit_order_column || !query_constraints_list.empty?)
-          # use query_constraints_list as the order clause if there is no implicit_order_column
-          # otherwise remove the implicit order column from the query constraints list if it's there
-          # and prepend it to the beginning of the list
-          order_columns = implicit_order_column.nil? ? query_constraints_list : ([implicit_order_column] | query_constraints_list)
-          order(*order_columns.map { |column| table[column].asc })
+        if order_values.empty? && (implicit_order_column || !query_constraints_list.nil? || primary_key)
+          order(_order_columns.map { |column| table[column].asc })
         else
           self
         end
       end
+
+      def _order_columns
+        oc = []
+
+        oc << implicit_order_column if implicit_order_column
+        oc << query_constraints_list if query_constraints_list
+
+        if primary_key && query_constraints_list.nil?
+          oc << primary_key
+        end
+
+        oc.flatten.uniq.compact
+      end
   end
 end

activerecord/test/cases/persistence_test.rb

@@ -1439,13 +1439,13 @@ def test_primary_key_stays_the_same
     assert_equal("id", ClothingItem.primary_key)
   end
 
-  def test_query_constraints_list_is_an_empty_array_if_primary_key_is_nil
+  def test_query_constraints_list_is_nil_if_primary_key_is_nil
     klass = Class.new(ActiveRecord::Base) do
       self.table_name = "developers_projects"
     end
 
     assert_nil klass.primary_key
-    assert_empty klass.query_constraints_list
+    assert_nil klass.query_constraints_list
   end
 
   def test_query_constraints_uses_primary_key_by_default

activerecord/test/models/sharded/blog_post.rb

@@ -6,6 +6,6 @@ class BlogPost < ActiveRecord::Base
     query_constraints :blog_id, :id
 
     belongs_to :blog
-    has_many :comments, foreign_key: [:blog_id, :blog_post_id]
+    has_many :comments, query_constraints: [:blog_id, :blog_post_id]
   end
 end

activerecord/test/models/sharded/comment.rb

@@ -5,7 +5,7 @@ class Comment < ActiveRecord::Base
     self.table_name = :sharded_comments
     query_constraints :blog_id, :id
 
-    belongs_to :blog_post, foreign_key: [:blog_id, :blog_post_id]
+    belongs_to :blog_post, query_constraints: [:blog_id, :blog_post_id]
     belongs_to :blog
   end
 end