id column rewriting

lib/has_aux_table/aux_table_config.rb

@@ -166,6 +166,7 @@ module HasAuxTable
               next
             end
           else
+            # TODO - add support for ActiveRecord::Reflection::ThroughReflection
             raise "unsupported association type: #{assoc.class}"
           end
         end

lib/has_aux_table/model_class_helper.rb

@@ -13,12 +13,12 @@ module HasAuxTable
       column_names.include?(name.to_s)
     end
 
-    sig { params(name: String).returns(T::Boolean) }
+    sig { params(name: T.any(String, Symbol)).returns(T::Boolean) }
     def is_primary_key?(name)
       primary_keys.include?(name.to_sym)
     end
 
-    sig { params(name: String).returns(T::Boolean) }
+    sig { params(name: T.any(String, Symbol)).returns(T::Boolean) }
     def is_type_key?(name)
       type_key == name.to_s
     end

lib/has_aux_table/relation_extensions.rb

@@ -51,18 +51,24 @@ class ActiveRecord::Associations::AssociationScope
       klass = refl.klass
       next unless klass.is_a?(HasAuxTable::ClassMethods)
       next unless aux_config = klass.aux_table_for(refl.join_primary_key)
-      aux_table = aux_config.aux.klass.table_name
+      aux_table =
-      main_table = aux_config.main.klass.table_name
+        scope.connection.quote_table_name(aux_config.aux.klass.table_name)
+      main_table =
+        scope.connection.quote_table_name(aux_config.main.klass.table_name)
       main_keys =
-        aux_config.main.primary_keys.map { |key| "'#{main_table}'.'#{key}'" }
+        aux_config.main.primary_keys.map do |key|
+          "#{main_table}.#{scope.connection.quote_column_name(key)}"
+        end
       aux_keys =
-        aux_config.aux.primary_keys.map { |key| "'#{aux_table}'.'#{key}'" }
+        aux_config.aux.primary_keys.map do |key|
+          "#{aux_table}.#{scope.connection.quote_column_name(key)}"
+        end
       join_clause =
         main_keys
           .zip(aux_keys)
           .map { |(main_key, aux_key)| "#{main_key} = #{aux_key}" }
           .join(" AND ")
-      scope.joins!("INNER JOIN '#{main_table}' ON (#{join_clause})")
+      scope.joins!("INNER JOIN #{main_table} ON (#{join_clause})")
     end if association.is_a?(
       ActiveRecord::Associations::HasManyThroughAssociation
     )
@@ -122,7 +128,26 @@ module HasAuxTable
       relation_class.send(:define_method, :pluck) do |*column_names|
         T.bind(self, ActiveRecord::Relation)
         if (predicates = Util.try_relation_optimization(self, aux_config))
-          aux_relation = aux_config.aux.klass.where(predicates)
+          aux_relation = aux_config.aux.klass.all
+          column_names.map! do |column_name|
+            if aux_config.main.is_primary_key?(column_name)
+              aux_table_name =
+                self.connection.quote_table_name(
+                  aux_config.aux.klass.table_name
+                )
+              pkey_column_name =
+                self.connection.quote_column_name(
+                  aux_config.aux.klass.primary_key
+                )
+              column_name = self.connection.quote_column_name(column_name)
+              Arel.sql(
+                "#{aux_table_name}.#{pkey_column_name} AS #{column_name}"
+              )
+            else
+              column_name
+            end
+          end
+          aux_relation.where!(predicates) if predicates.any?
           aux_relation.pluck(*column_names)
         else
           pluck_method.bind(self).call(*column_names)
@@ -133,10 +158,15 @@ module HasAuxTable
       relation_class.send(
         :define_method,
         :calculate
-      ) do |operation, column_name|
+      ) do |operation, column_name = nil|
         T.bind(self, ActiveRecord::Relation)
         if (predicates = Util.try_relation_optimization(self, aux_config))
-          aux_relation = aux_config.aux.klass.where(predicates)
+          aux_relation = aux_config.aux.klass.all
+
+          if column_name && aux_config.main.is_primary_key?(column_name)
+            column_name = aux_config.aux.klass.primary_key
+          end
+          aux_relation.where!(predicates) if predicates.any?
           aux_relation.calculate(operation, column_name)
         else
           calculate_method.bind(self).call(operation, column_name)

spec/has_aux_table_spec.rb

@@ -235,7 +235,7 @@ RSpec.describe HasAuxTable do
     it "allows saving the model with auxiliary columns" do
       car = Car.create!(name: "Honda Civic")
       queries =
-        SpecHelper.capture_queries do
+        capture_queries do
           car.fuel_type = "hybrid"
           car.engine_size = 1.8
           car.save!
@@ -424,7 +424,7 @@ RSpec.describe HasAuxTable do
 
         it "loads single model with auxiliary data in one query using find" do
           queries =
-            SpecHelper.capture_queries do
+            capture_queries do
               car = Car.find(@car1.id)
               # Access auxiliary attributes to ensure they're loaded
               car.fuel_type
@@ -436,7 +436,7 @@ RSpec.describe HasAuxTable do
 
         it "loads single model with auxiliary data in one query using find_by" do
           queries =
-            SpecHelper.capture_queries do
+            capture_queries do
               car = Car.find_by(name: "Toyota Prius")
               # Access auxiliary attributes to ensure they're loaded
               car.fuel_type
@@ -448,7 +448,7 @@ RSpec.describe HasAuxTable do
 
         it "loads multiple models with auxiliary data in one query using where" do
           queries =
-            SpecHelper.capture_queries do
+            capture_queries do
               cars = Car.where(fuel_type: %w[hybrid electric])
               # Access auxiliary attributes for all cars
               cars.each do |car|
@@ -473,7 +473,7 @@ RSpec.describe HasAuxTable do
 
           cars = nil
           queries =
-            SpecHelper.capture_queries do
+            capture_queries do
               cars = Car.where(fuel_type: "gasoline")
               # Access auxiliary attributes for all cars - should not trigger additional queries
               cars.each do |car|
@@ -489,7 +489,7 @@ RSpec.describe HasAuxTable do
 
         it "uses single query when ordering by auxiliary columns" do
           queries =
-            SpecHelper.capture_queries do
+            capture_queries do
               cars = Car.where(engine_size: 1.0..3.0).order(:engine_size)
               # Access all attributes
               cars.each do |car|
@@ -504,7 +504,7 @@ RSpec.describe HasAuxTable do
 
         it "uses single query for complex auxiliary column queries" do
           queries =
-            SpecHelper.capture_queries do
+            capture_queries do
               cars =
                 Car.where(fuel_type: "hybrid").or(Car.where(engine_size: 0.0))
               # Access all attributes
@@ -520,7 +520,7 @@ RSpec.describe HasAuxTable do
 
         it "uses single query when finding by auxiliary columns" do
           queries =
-            SpecHelper.capture_queries do
+            capture_queries do
               car = Car.find_by(fuel_type: "hybrid", name: "Toyota Prius")
               # Access all attributes
               car.name
@@ -537,7 +537,7 @@ RSpec.describe HasAuxTable do
 
           # Now count queries when accessing auxiliary attributes
           queries =
-            SpecHelper.capture_queries do
+            capture_queries do
               car.fuel_type
               car.engine_size
               car.fuel_type? # presence check
@@ -550,7 +550,7 @@ RSpec.describe HasAuxTable do
 
         it "handles mixed queries with main and auxiliary columns in single query" do
           queries =
-            SpecHelper.capture_queries do
+            capture_queries do
               cars = Car.where(name: "Toyota Prius", fuel_type: "hybrid")
               cars.each do |car|
                 car.name
@@ -564,7 +564,7 @@ RSpec.describe HasAuxTable do
 
         it "uses single query for range queries on auxiliary columns" do
           queries =
-            SpecHelper.capture_queries do
+            capture_queries do
               cars = Car.where(engine_size: 0.0..1.9)
               cars.each do |car|
                 car.name
@@ -578,7 +578,7 @@ RSpec.describe HasAuxTable do
 
         it "maintains single query performance with limit and offset" do
           queries =
-            SpecHelper.capture_queries do
+            capture_queries do
               car = Car.where(fuel_type: %w[hybrid electric]).limit(1).first
               # Access auxiliary attributes
               car.fuel_type
@@ -983,7 +983,7 @@ RSpec.describe HasAuxTable do
       end
 
       it "reloads with one query" do
-        queries = SpecHelper.capture_queries { @car.reload }
+        queries = capture_queries { @car.reload }
         expect(queries.length).to eq(1)
       end
 
@@ -1027,16 +1027,12 @@ RSpec.describe HasAuxTable do
     expect(Car.count).to eq(1)
     expect(Boat.count).to eq(1)
 
-    expect(
+    expect(capture_queries { car = Vehicle.find(car.id) }.length).to eq(1)
-      SpecHelper.capture_queries { car = Vehicle.find(car.id) }.length
-    ).to eq(1)
     expect(car.fuel_type).to eq("gasoline")
     expect(car.engine_size).to eq(2.0)
     expect(car.name).to eq("Honda Civic")
 
-    expect(
+    expect(capture_queries { boat = Vehicle.find(boat.id) }.length).to eq(1)
-      SpecHelper.capture_queries { boat = Vehicle.find(boat.id) }.length
-    ).to eq(1)
     expect(boat.only_freshwater).to eq(true)
   end
 

spec/loading_optimizations_spec.rb

@@ -17,84 +17,105 @@ RSpec.describe "loading optimizations" do
 
     shared_examples "queries only the aux table" do
       it "queries only the aux table" do
-        expect(@queries.length).to eq(1)
+        expect(queries.length).to eq(1)
-        expect(@queries.first).not_to include("JOIN")
+        expect(queries.first).not_to include("JOIN")
-        expect(@queries.first).to match(/\bvehicles_car_aux\b/)
+        expect(queries.first).to match(/\bvehicles_car_aux\b/)
-        expect(@queries.first).not_to match(/\bvehicles\b/)
+        expect(queries.first).not_to match(/\bvehicles\b/)
       end
     end
 
     shared_examples "queries both tables" do
       it "queries both tables" do
-        expect(@queries.length).to eq(1)
+        expect(queries.length).to eq(1)
-        expect(@queries.first).to include("JOIN")
+        expect(queries.first).to include("JOIN")
-        expect(@queries.first).to match(/\bvehicles\b/)
+        expect(queries.first).to match(/\bvehicles\b/)
-        expect(@queries.first).to match(/\bvehicles_car_aux\b/)
+        expect(queries.first).to match(/\bvehicles_car_aux\b/)
       end
     end
 
     describe "pluck" do
       context "aux columns are referenced" do
-        before do
+        let_and_capture(:queries) { Car.pluck(:fuel_type) }
-          @queries =
+
-            SpecHelper.capture_queries do
+        it "returns the correct result" do
-              expect(Car.pluck(:fuel_type)).to eq(%w[gasoline hybrid electric])
+          expect(queries.result).to eq(%w[gasoline hybrid electric])
-            end
         end
 
         it_behaves_like "queries only the aux table"
       end
 
       context "aux columns are chained on a where clause" do
-        before do
+        let_and_capture(:queries) do
-          @queries =
+          Car.where(engine_size: 1.4..1.9).pluck(:fuel_type)
-            SpecHelper.capture_queries do
+        end
-              expect(Car.where(engine_size: 1.4..1.9).pluck(:fuel_type)).to eq(
+
-                %w[hybrid electric]
+        it "returns the correct result" do
-              )
+          expect(queries.result).to eq(%w[hybrid electric])
-            end
         end
 
         it_behaves_like "queries only the aux table"
         it "applies the BETWEEN clause" do
-          expect(@queries.first).to include("BETWEEN")
+          expect(queries.first).to include("BETWEEN")
         end
       end
 
       context "main table columns are referenced" do
-        before do
+        let(:queries) do
-          @queries =
+          capture_queries { Car.where(name: "Toyota Camry").pluck(:fuel_type) }
-            SpecHelper.capture_queries do
+        end
-              expect(Car.where(name: "Toyota Camry").pluck(:fuel_type)).to eq(
+
-                %w[gasoline]
+        it "returns the correct result" do
-              )
+          expect(queries.result).to eq(%w[gasoline])
-            end
         end
 
         it_behaves_like "queries both tables"
       end
 
       context "main table columns are chained on a where clause" do
-        before do
+        let(:queries) do
-          @queries =
+          capture_queries { Car.where(name: "Toyota Camry").pluck(:fuel_type) }
-            SpecHelper.capture_queries do
+        end
-              expect(Car.where(name: "Toyota Camry").pluck(:fuel_type)).to eq(
+
-                %w[gasoline]
+        it "returns the correct result" do
-              )
+          expect(queries.result).to eq(%w[gasoline])
-            end
         end
 
         it_behaves_like "queries both tables"
       end
 
       context "multiple columns" do
-        before do
+        let(:queries) do
-          @queries =
+          capture_queries { Car.pluck(:fuel_type, :engine_size) }
-            SpecHelper.capture_queries do
+        end
-              expect(Car.pluck(:fuel_type, :engine_size)).to eq(
+
-                [["gasoline", 2.0], ["hybrid", 1.5], ["electric", 1.8]]
+        it "returns the correct result" do
-              )
+          expect(queries.result).to eq(
-            end
+            [["gasoline", 2.0], ["hybrid", 1.5], ["electric", 1.8]]
+          )
+        end
+
+        it_behaves_like "queries only the aux table"
+      end
+
+      context "multiple columns with a where clause" do
+        let(:queries) do
+          capture_queries do
+            Car.where(name: "Toyota Camry").pluck(:fuel_type, :engine_size)
+          end
+        end
+
+        it "returns the correct result" do
+          expect(queries.result).to eq([["gasoline", 2.0]])
+        end
+
+        it_behaves_like "queries both tables"
+      end
+
+      context "querying the id column" do
+        let_and_capture(:queries) { Car.pluck(:id) }
+
+        it "renames the base_table_id to id" do
+          expect(queries.first).to include("\"base_table_id\" AS \"id\"")
         end
 
         it_behaves_like "queries only the aux table"
@@ -103,50 +124,56 @@ RSpec.describe "loading optimizations" do
 
     describe "maximum" do
       context "aux columns are referenced" do
-        before do
+        let_and_capture(:queries) { Car.maximum(:engine_size) }
-          @queries =
-            SpecHelper.capture_queries do
-              expect(Car.maximum(:engine_size)).to eq(2.0)
-            end
-        end
 
         it_behaves_like "queries only the aux table"
       end
 
       context "aux columns are chained on a where clause" do
-        before do
+        let_and_capture(:queries) do
-          @queries =
+          Car.where(engine_size: 1.4..1.9).maximum(:engine_size)
-            SpecHelper.capture_queries do
+        end
-              expect(
+
-                Car.where(engine_size: 1.4..1.9).maximum(:engine_size)
+        it "returns the correct result" do
-              ).to eq(1.8)
+          expect(queries.result).to eq(1.8)
-            end
         end
 
         it_behaves_like "queries only the aux table"
       end
 
       context "main table columns are referenced" do
-        before do
+        let_and_capture(:queries) do
-          @queries =
+          Car.where(name: "Toyota Camry").maximum(:engine_size)
-            SpecHelper.capture_queries do
+        end
-              expect(
+
-                Car.where(name: "Toyota Camry").maximum(:engine_size)
+        it "returns the correct result" do
-              ).to eq(2.0)
+          expect(queries.result).to eq(2.0)
-            end
         end
 
         it_behaves_like "queries both tables"
       end
 
+      context "id column is referenced" do
+        let_and_capture(:queries) { Car.maximum(:id) }
+
+        it "renames the base_table_id to id" do
+          expect(queries.first).to include('"base_table_id"')
+        end
+
+        it "has the right result" do
+          expect(queries.result).to eq(3)
+        end
+
+        it_behaves_like "queries only the aux table"
+      end
+
       context "main table columns are chained on a where clause" do
-        before do
+        let_and_capture(:queries) do
-          @queries =
+          Car.where(name: "Toyota Camry").maximum(:engine_size)
-            SpecHelper.capture_queries do
+        end
-              expect(
+
-                Car.where(name: "Toyota Camry").maximum(:engine_size)
+        it "returns the correct result" do
-              ).to eq(2.0)
+          expect(queries.result).to eq(2.0)
-            end
         end
 
         it_behaves_like "queries both tables"

spec/spec_helper.rb

@@ -54,6 +54,36 @@ RSpec.configure do |config|
       raise ActiveRecord::Rollback
     end
   end
+
+  config.include(
+    Module.new do
+      extend T::Sig
+      sig do
+        type_parameters(:T)
+          .params(block: T.proc.returns(T.type_parameter(:T)))
+          .returns(SpecHelper::CaptureQueries[T.type_parameter(:T)])
+      end
+      def capture_queries(&block)
+        SpecHelper.capture_queries(&block)
+      end
+    end
+  )
+
+  config.extend(
+    Module.new do
+      extend T::Sig
+
+      sig do
+        type_parameters(:T)
+          .params(binding: Symbol, block: T.proc.returns(T.type_parameter(:T)))
+          .returns(T.type_parameter(:T))
+      end
+      def let_and_capture(binding, &block)
+        T.bind(self, RSpec::Core::MemoizedHelpers::ClassMethods)
+        let(binding) { SpecHelper.capture_queries(&block) }
+      end
+    end
+  )
 end
 
 module SpecHelper
@@ -61,8 +91,28 @@ module SpecHelper
   extend T::Helpers
   LOG_QUERIES = T.let(false, T::Boolean)
 
+  class CaptureQueries < Array
+    extend T::Sig
+    extend T::Generic
+    Elem = type_member { { fixed: String } }
+    Result = type_member
+
+    sig { params(queries: T::Array[String], result: Result).void }
+    def initialize(queries, result)
+      super(queries)
+      @result = result
+    end
+
+    sig { returns(Result) }
+    attr_reader :result
+  end
+
   # Helper method to count queries
-  sig { params(block: T.proc.void).returns(T::Array[String]) }
+  sig do
+    type_parameters(:T)
+      .params(block: T.proc.returns(T.type_parameter(:T)))
+      .returns(CaptureQueries[T.type_parameter(:T)])
+  end
   def self.capture_queries(&block)
     queries = T.let([], T::Array[String])
     query_callback =
@@ -70,6 +120,7 @@ module SpecHelper
         queries << values[:sql]
       end
 
+    result = T.let(nil, T.untyped)
     ActiveSupport::Notifications.subscribed(
       query_callback,
       "sql.active_record"
@@ -79,10 +130,11 @@ module SpecHelper
         ActiveRecord::Base.logger = Logger.new(STDOUT)
         ActiveRecord::Base.logger.level = Logger::DEBUG
       end
-      block.call
+      result = block.call
+    ensure
       ActiveRecord::Base.logger = old_logger if LOG_QUERIES
     end
 
-    queries
+    CaptureQueries.new(queries, result)
   end
 end