Refactor association handling in `PredicateBuilder`

I'm attempting to remove `klass` as a dependency of the predicate
builder, in favor of an object that better represents what we're using
it for. The only part of this which doesn't fit nicely into that picture
is the check for an association being polymorphic. Since I'm not yet
sure what that is going to look like, I've moved this logic into another
class in an attempt to separate things that will change from things that
won't.

2 files changed, 61 insertions, 21 deletions

diff --git a/activerecord/lib/active_record/relation/predicate_builder.rb b/activerecord/lib/active_record/relation/predicate_builder.rb
index 61db27d150..b0a4f27049 100644
--- a/activerecord/lib/active_record/relation/predicate_builder.rb
+++ b/activerecord/lib/active_record/relation/predicate_builder.rb
@@ -1,6 +1,7 @@
 module ActiveRecord
   class PredicateBuilder # :nodoc:
     require 'active_record/relation/predicate_builder/array_handler'
+    require 'active_record/relation/predicate_builder/association_query_handler'
     require 'active_record/relation/predicate_builder/base_handler'
     require 'active_record/relation/predicate_builder/basic_object_handler'
     require 'active_record/relation/predicate_builder/class_handler'
@@ -18,6 +19,7 @@ module ActiveRecord
       register_handler(Range, RangeHandler.new)
       register_handler(Relation, RelationHandler.new)
       register_handler(Array, ArrayHandler.new(self))
+      register_handler(AssociationQueryValue, AssociationQueryHandler.new(self))
     end
 
     def resolve_column_aliases(hash)
@@ -36,35 +38,16 @@ module ActiveRecord
     end
 
     def expand(column, value)
-      queries = []
-
       # Find the foreign key when using queries such as:
       # Post.where(author: author)
       #
       # For polymorphic relationships, find the foreign key and type:
       # PriceEstimate.where(estimate_of: treasure)
       if klass && reflection = klass._reflect_on_association(column)
-        if reflection.polymorphic? && base_class = polymorphic_base_class_from_value(value)
-          queries << build(table[reflection.foreign_type], base_class.name)
-        end
-
-        column = reflection.foreign_key
+        value = AssociationQueryValue.new(reflection, value)
       end
 
-      queries << build(table[column], value)
-      queries
-    end
-
-    def polymorphic_base_class_from_value(value)
-      case value
-      when Relation
-        value.klass.base_class
-      when Array
-        val = value.compact.first
-        val.class.base_class if val.is_a?(Base)
-      when Base
-        value.class.base_class
-      end
+      build(table[column], value)
     end
 
     def self.references(attributes)
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
new file mode 100644
index 0000000000..bda2c02a10
--- /dev/null
+++ b/activerecord/lib/active_record/relation/predicate_builder/association_query_handler.rb
@@ -0,0 +1,57 @@
+module ActiveRecord
+  class PredicateBuilder
+    class AssociationQueryHandler # :nodoc:
+      def initialize(predicate_builder)
+        @predicate_builder = predicate_builder
+      end
+
+      def call(attribute, value)
+        queries = {}
+
+        if value.base_class
+          queries[value.association.foreign_type] = value.base_class.name
+        end
+
+        queries[value.association.foreign_key] = value.ids
+        predicate_builder.build_from_hash(queries)
+      end
+
+      protected
+
+      attr_reader :predicate_builder
+    end
+
+    class AssociationQueryValue # :nodoc:
+      attr_reader :association, :value
+
+      def initialize(association, value)
+        @association = association
+        @value = value
+      end
+
+      def ids
+        value
+      end
+
+      def base_class
+        if association.polymorphic?
+          @base_class ||= polymorphic_base_class_from_value
+        end
+      end
+
+      private
+
+      def polymorphic_base_class_from_value
+        case value
+        when Relation
+          value.klass.base_class
+        when Array
+          val = value.compact.first
+          val.class.base_class if val.is_a?(Base)
+        when Base
+          value.class.base_class
+        end
+      end
+    end
+  end
+end