diff --git a/lib/has_aux_table/relation_extensions.rb b/lib/has_aux_table/relation_extensions.rb index 89a1d87..a109ce2 100644 --- a/lib/has_aux_table/relation_extensions.rb +++ b/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 diff --git a/lib/has_aux_table/util.rb b/lib/has_aux_table/util.rb index 6767d93..fa1e86a 100644 --- a/lib/has_aux_table/util.rb +++ b/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 diff --git a/spec/loading_optimizations_spec.rb b/spec/loading_optimizations_spec.rb index 6564a8a..8ce3c73 100644 --- a/spec/loading_optimizations_spec.rb +++ b/spec/loading_optimizations_spec.rb @@ -6,12 +6,16 @@ require "spec_helper" RSpec.describe "loading optimizations" do context "cars table" do before do - Car.create!(name: "Toyota Camry", fuel_type: "gasoline") - Car.create!(name: "Toyota Prius", fuel_type: "hybrid") - Car.create!(name: "Toyota Corolla", fuel_type: "electric") + Car.create!(name: "Toyota Camry", fuel_type: "gasoline", engine_size: 2.0) + Car.create!(name: "Toyota Prius", fuel_type: "hybrid", engine_size: 1.5) + Car.create!( + name: "Toyota Corolla", + fuel_type: "electric", + engine_size: 1.8 + ) end - it "queries only the aux table if plucking values that are on the aux table" do + it "queries only the aux table if no main table columns are referenced" do queries = SpecHelper.capture_queries do expect(Car.pluck(:fuel_type)).to eq(%w[gasoline hybrid electric]) @@ -21,6 +25,21 @@ RSpec.describe "loading optimizations" do expect(queries.first).not_to include("JOIN") end + it "queries only the aux table if all columns are on the aux table" do + queries = + SpecHelper.capture_queries do + expect(Car.where(engine_size: 1.4..1.9).pluck(:fuel_type)).to eq( + %w[hybrid electric] + ) + end + + expect(queries.length).to eq(1) + expect(queries.first).not_to include("JOIN") + expect(queries.first).to include("BETWEEN") + expect(queries.first).to match(/\bvehicles_car_aux\b/) + expect(queries.first).not_to match(/\bvehicles\b/) + end + it "queries both tables if main table column is referenced" do queries = SpecHelper.capture_queries do