Merge pull request #28713 from kamipo/expose_queries_for_association_queries

Expose `queries` for `AssociationQueryValue` and `PolymorphicArrayValue`

2 files changed, 44 insertions, 48 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..9141d9c537 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,15 +20,7 @@ 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
-        predicate_builder.build_from_hash(queries)
+        predicate_builder.build_from_hash(value.queries)
       end
 
       # TODO Change this to private once we've dropped Ruby 2.2 support.
@@ -43,38 +38,26 @@ module ActiveRecord
         @value = value
       end
 
-      def ids
-        case value
-        when Relation
-          value.select(primary_key)
-        when Array
-          value.map { |v| convert_to_id(v) }
-        else
-          convert_to_id(value)
-        end
-      end
-
-      def base_class
-        if associated_table.polymorphic_association?
-          @base_class ||= polymorphic_base_class_from_value
-        end
+      def queries
+        { associated_table.association_foreign_key.to_s => ids }
       end
 
       private
-
-        def primary_key
-          associated_table.association_primary_key(base_class)
-        end
-
-        def polymorphic_base_class_from_value
+        def ids
           case value
           when Relation
-            value.klass.base_class
-          when Base
-            value.class.base_class
+            value.select_values.empty? ? value.select(primary_key) : value
+          when Array
+            value.map { |v| convert_to_id(v) }
+          else
+            convert_to_id(value)
           end
         end
 
+        def primary_key
+          associated_table.association_primary_key
+        end
+
         def convert_to_id(value)
           case value
           when Base
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..c2f136256b 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
@@ -6,12 +6,7 @@ module ActiveRecord
       end
 
       def call(attribute, value)
-        table = value.associated_table
-        queries = value.type_to_ids_mapping.map do |type, ids|
-          { table.association_foreign_type.to_s => type, table.association_foreign_key.to_s => ids }
-        end
-
-        predicates = queries.map { |query| predicate_builder.build_from_hash(query) }
+        predicates = value.queries.map { |query| predicate_builder.build_from_hash(query) }
 
         if predicates.size > 1
           type_and_ids_predicates = predicates.map { |type_predicate, id_predicate| Arel::Nodes::Grouping.new(type_predicate.and(id_predicate)) }
@@ -36,23 +31,41 @@ module ActiveRecord
         @values = values
       end
 
-      def type_to_ids_mapping
-        default_hash = Hash.new { |hsh, key| hsh[key] = [] }
-        values.each_with_object(default_hash) { |value, hash| hash[base_class(value).name] << convert_to_id(value) }
+      def queries
+        type_to_ids_mapping.map do |type, ids|
+          {
+            associated_table.association_foreign_type.to_s => type,
+            associated_table.association_foreign_key.to_s => ids
+          }
+        end
       end
 
       private
+        def type_to_ids_mapping
+          default_hash = Hash.new { |hsh, key| hsh[key] = [] }
+          values.each_with_object(default_hash) { |value, hash| hash[base_class(value).name] << convert_to_id(value) }
+        end
 
         def primary_key(value)
           associated_table.association_primary_key(base_class(value))
         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