optimize pluck

lib/has_aux_table/relation_extensions.rb

@@ -122,50 +122,42 @@ module HasAuxTable
       relation_class.send(:define_method, :pluck) do |column_names|
         T.bind(self, ActiveRecord::Relation)
 
-        filtered_attributes =
-          self.where_clause.extract_attributes.select do |attr|
-            column_name = attr.name
-            if attr.relation == aux_config.main.table
-              # if it's on the main table, ignore if it if's the primary key or type key
-              next false if aux_config.main.is_primary_key?(column_name)
-              next false if aux_config.main.is_type_key?(column_name)
+        all_predicates =
+          self.where_clause.instance_eval do
+            T.cast(predicates, T::Array[Arel::Nodes::Binary])
+          end
+
+        filtered_predicates =
+          all_predicates.filter do |node|
+            if node.is_a?(Arel::Nodes::Equality)
+              if Util.is_same_table?(node.left.relation, aux_config.main.table)
+                # if it's on the main table, ignore if it if's the primary key or type key
+                name = node.left.name
+                next false if aux_config.main.is_primary_key?(name)
+                next false if aux_config.main.is_type_key?(name)
+              end
             end
             true
           end
 
         all_on_aux_table =
-          filtered_attributes.all? do |attr|
-            column_name = attr.name
-
+          filtered_predicates.all? do |node|
             # if it's a field on the aux table, then it can be plucked
-            if attr.relation == aux_config.aux.table
-              puts "optimize: #{column_name} is on #{attr.relation.name}"
-              next true
-            end
-
-            # if it's on the main table, ignore if it if's the primary key or type key
-            if attr.relation == aux_config.main.table
-              if aux_config.main.is_primary_key?(column_name)
-                puts "optimize: #{column_name} is primary key on #{aux_config.main.table.name}"
-                next true
-              end
-              if aux_config.main.is_type_key?(column_name)
-                puts "optimize: #{column_name} is type key on #{aux_config.main.table.name}"
-                next true
-              end
-            end
-
-            puts "skip optimization: #{column_name} is on #{attr.relation.name}"
-            false
+            Util.is_same_table?(node.left.relation, aux_config.aux.table)
           end
 
         if all_on_aux_table
-          Kernel.puts "pluck is only for aux columns: #{column_names}"
-          binding.pry
-          aux_relation = aux_config.aux.klass.where(where_clause)
+          # the eager load generates a join which creates table alias nodes on attributes instead
+          # of the original table, so we need to replace those with the original table
+          filtered_predicates.each do |node|
+            if (attribute = node.left) && (table_alias = attribute.relation) &&
+                 table_alias.is_a?(Arel::Nodes::TableAlias)
+              attribute.relation = table_alias.left
+            end
+          end
+          aux_relation = aux_config.aux.klass.where(filtered_predicates)
           aux_relation.pluck(*column_names)
         else
-          Kernel.puts "pluck proxied to original: #{column_names}"
           pluck_method.bind(self).call(*column_names)
         end
       end

lib/has_aux_table/util.rb

@@ -69,5 +69,13 @@ module HasAuxTable
         Kernel.raise("#{instance.class.name} not a #{klass.name}")
       end
     end
+
+    TableOrAlias = T.type_alias { T.any(Arel::Nodes::Node, Arel::Table) }
+    sig { params(left: TableOrAlias, right: TableOrAlias).returns(T::Boolean) }
+    def self.is_same_table?(left, right)
+      left_table = left.is_a?(Arel::Nodes::TableAlias) ? left.left : left
+      right_table = right.is_a?(Arel::Nodes::TableAlias) ? right.left : right
+      left_table == right_table
+    end
   end
 end