Extend ActiveRecord PredicateBuilder to support querying by polymorphic integer relations

Author: tgurgul

lib/polymorphic_integer_type.rb

@@ -1,5 +1,6 @@
 require "polymorphic_integer_type/version"
 require "polymorphic_integer_type/extensions"
 require "polymorphic_integer_type/mapping"
+require "polymorphic_integer_type/predicate_builder_extension"
 
 module PolymorphicIntegerType; end

lib/polymorphic_integer_type/extensions.rb

@@ -12,6 +12,11 @@ def belongs_to(name, scope = nil, options = {})
           foreign_type = reflections[name.to_s].foreign_type
           self._polymorphic_foreign_types << foreign_type
 
+          # Required way to dynamically define a class method on the model
+          singleton_class.__send__(:define_method, "#{foreign_type}_mapping") do
+            mapping
+          end
+
           define_method foreign_type do
             t = super()
             mapping[t]

lib/polymorphic_integer_type/predicate_builder_extension.rb

@@ -0,0 +1,36 @@
+module PolymorphicIntegerType
+  module PredicateBuilderExtension
+    def self.expand(klass, table, 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)
+        base_class = polymorphic_base_class_from_value(value)
+
+        if reflection.polymorphic? && base_class
+          if klass.respond_to?("#{column}_type_mapping")
+            queries << build(table[reflection.foreign_type], klass.send("#{column}_type_mapping").key(base_class.to_s))
+          else
+            queries << build(table[reflection.foreign_type], base_class)
+          end
+        end
+
+        column = reflection.foreign_key
+
+        if base_class
+          primary_key = reflection.association_primary_key(base_class)
+          value = convert_value_to_association_ids(value, primary_key)
+        end
+      end
+
+      queries << build(table[column], value)
+      queries
+    end
+  end
+end
+
+ActiveRecord::PredicateBuilder.prepend(PolymorphicIntegerType::PredicateBuilderExtension)