Fix bug in Rails 5.0 when querying polymorphic type columns

The existing patch on PolymorphicArrayValue only handled queries
that passed multiple values to a polymorphic type column, for example:

  Link.where(source: [source_1, source_2])

In order to handle passing a single value, we need to patch the
AssociationQueryHandler.

This issue only exists in Rails 5.0 and 5.1; it's been fixed in 5.2.

Author: mctaylorpants

lib/polymorphic_integer_type.rb

@@ -1,10 +1,17 @@
+ACTIVE_RECORD_VERSION = Gem::Version.new(ActiveRecord::VERSION::STRING)
+
 require "polymorphic_integer_type/version"
 require "polymorphic_integer_type/extensions"
 require "polymorphic_integer_type/mapping"
-if Gem::Version.new(ActiveRecord::VERSION::STRING) < Gem::Version.new("5")
-  require "polymorphic_integer_type/predicate_builder_extension"
+
+if ACTIVE_RECORD_VERSION < Gem::Version.new("5")
+  require "polymorphic_integer_type/activerecord_4/predicate_builder_extension"
 else
-  require "polymorphic_integer_type/polymorphic_array_value_extension"
+  require "polymorphic_integer_type/activerecord_5_0_0/polymorphic_array_value_extension"
+end
+
+if ACTIVE_RECORD_VERSION >= Gem::Version.new("5.0") && ACTIVE_RECORD_VERSION < Gem::Version.new("5.2.0")
+  require "polymorphic_integer_type/activerecord_5_0_0/association_query_handler_extension"
 end
 
 module PolymorphicIntegerType; end

…eger_type/predicate_builder_extension.rb …erecord_4/predicate_builder_extension.rblib/polymorphic_integer_type/predicate_builder_extension.rb renamed to lib/polymorphic_integer_type/activerecord_4/predicate_builder_extension.rb lib/polymorphic_integer_type/predicate_builder_extension.rb renamed to lib/polymorphic_integer_type/activerecord_4/predicate_builder_extension.rb

@@ -0,0 +1,34 @@
+module PolymorphicIntegerType
+  module AssociationQueryHandlerExtension
+    def call(attribute, value)
+      queries = {}
+      table = value.associated_table
+
+      if value.base_class
+        queries[table.association_foreign_type.to_s] = polymorphic_value_for(value)
+      end
+
+      queries[table.association_foreign_key.to_s] = value.ids
+      predicate_builder.build_from_hash(queries)
+    end
+
+    protected
+
+    def polymorphic_value_for(value)
+      table = value.associated_table
+      association = table.send(:association)
+      klass = association.active_record
+      name = association.name
+
+      if klass.respond_to?("#{name}_type_mapping")
+        klass.send("#{name}_type_mapping").key(value.base_class.sti_name)
+      else
+        value.base_class.name
+      end
+    end
+
+
+  end
+end
+
+ActiveRecord::PredicateBuilder::AssociationQueryHandler.prepend(PolymorphicIntegerType::AssociationQueryHandlerExtension)

lib/polymorphic_integer_type/activerecord_5_0_0/association_query_handler_extension.rb

@@ -34,9 +34,15 @@
   context "when querying the associations" do
     let(:source) { cat }
     let(:target) { nil }
+
     it "properly finds the object with a where" do
       expect(Link.where(source: source, id: link.id).first).to eql link
     end
+
+    it "properly finds the object when passing an array of sources" do
+      expect(Link.where(source: [source])).to eq [link]
+    end
+
     it "properly finds the object with a find_by" do
       expect(Link.find_by(source: source, id: link.id)).to eql link
     end