Avoid to handle polymorphic association for `AssociationQueryHandler`

It should be handled by `PolymorphicArrayHandler` if polymorphic association.

2 files changed, 21 insertions, 29 deletions

diff --git a/activerecord/lib/active_record/relation/predicate_builder/association_query_handler.rb b/activerecord/lib/active_record/relation/predicate_builder/association_query_handler.rb
index 29860ec677..82da4689f0 100644
--- a/activerecord/lib/active_record/relation/predicate_builder/association_query_handler.rb
+++ b/activerecord/lib/active_record/relation/predicate_builder/association_query_handler.rb
@@ -3,12 +3,15 @@ module ActiveRecord
     class AssociationQueryHandler # :nodoc:
       def self.value_for(table, column, value)
         associated_table = table.associated_table(column)
-        klass = if associated_table.polymorphic_association? && ::Array === value && value.first.is_a?(Base)
-          PolymorphicArrayValue
-        else
-          AssociationQueryValue
+        if associated_table.polymorphic_association?
+          case value.is_a?(Array) ? value.first : value
+          when Base, Relation
+            value = [value] unless value.is_a?(Array)
+            klass = PolymorphicArrayValue
+          end
         end
 
+        klass ||= AssociationQueryValue
         klass.new(associated_table, value)
       end
 
@@ -17,14 +20,8 @@ module ActiveRecord
       end
 
       def call(attribute, value)
-        queries = {}
-
         table = value.associated_table
-        if value.base_class
-          queries[table.association_foreign_type.to_s] = value.base_class.name
-        end
-
-        queries[table.association_foreign_key.to_s] = value.ids
+        queries = { table.association_foreign_key.to_s => value.ids }
         predicate_builder.build_from_hash(queries)
       end
 
@@ -54,25 +51,10 @@ module ActiveRecord
         end
       end
 
-      def base_class
-        if associated_table.polymorphic_association?
-          @base_class ||= polymorphic_base_class_from_value
-        end
-      end
-
       private
 
         def primary_key
-          associated_table.association_primary_key(base_class)
-        end
-
-        def polymorphic_base_class_from_value
-          case value
-          when Relation
-            value.klass.base_class
-          when Base
-            value.class.base_class
-          end
+          associated_table.association_primary_key
         end
 
         def convert_to_id(value)
diff --git a/activerecord/lib/active_record/relation/predicate_builder/polymorphic_array_handler.rb b/activerecord/lib/active_record/relation/predicate_builder/polymorphic_array_handler.rb
index 335124c952..21cc324e5e 100644
--- a/activerecord/lib/active_record/relation/predicate_builder/polymorphic_array_handler.rb
+++ b/activerecord/lib/active_record/relation/predicate_builder/polymorphic_array_handler.rb
@@ -48,11 +48,21 @@ module ActiveRecord
         end
 
         def base_class(value)
-          value.class.base_class
+          case value
+          when Base
+            value.class.base_class
+          when Relation
+            value.klass.base_class
+          end
         end
 
         def convert_to_id(value)
-          value._read_attribute(primary_key(value))
+          case value
+          when Base
+            value._read_attribute(primary_key(value))
+          when Relation
+            value.select(primary_key(value))
+          end
         end
     end
   end