[wip] loading optimizations

lib/has_aux_table/model_class_helper.rb

@@ -13,6 +13,26 @@ module HasAuxTable
       column_names.include?(name.to_s)
     end
 
+    sig { params(name: String).returns(T::Boolean) }
+    def is_primary_key?(name)
+      primary_keys.include?(name.to_sym)
+    end
+
+    sig { params(name: String).returns(T::Boolean) }
+    def is_type_key?(name)
+      type_key == name.to_s
+    end
+
+    sig { returns(T.nilable(String)) }
+    def type_key
+      self.klass.inheritance_column
+    end
+
+    sig { returns(Arel::Table) }
+    def table
+      self.klass.arel_table
+    end
+
     sig { returns(T::Array[Symbol]) }
     def primary_keys
       @primary_keys ||=

lib/has_aux_table/relation_extensions.rb

@@ -118,6 +118,58 @@ module HasAuxTable
           ActiveRecord::Associations::CollectionProxy
         )
 
+      pluck_method = relation_class.instance_method(:pluck)
+      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)
+            end
+            true
+          end
+
+        all_on_aux_table =
+          filtered_attributes.all? do |attr|
+            column_name = attr.name
+
+            # 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
+          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)
+          aux_relation.pluck(*column_names)
+        else
+          Kernel.puts "pluck proxied to original: #{column_names}"
+          pluck_method.bind(self).call(*column_names)
+        end
+      end
+
       [
         [relation_class, :build_where_clause],
         [collection_proxy_class, :where]

lib/has_aux_table/util.rb

@@ -31,8 +31,12 @@ module HasAuxTable
         )
 
       target.send(define_method, method_name) do |*args, **kwargs, &block|
-        method = is_instance_method ? target_method.bind(self) : target_method
+        if is_instance_method
-        T.unsafe(hook_block).call(method, *args, **kwargs, &block)
+          method = target_method.bind(self)
+          T.unsafe(hook_block).call(method, *args, **kwargs, &block)
+        else
+          T.unsafe(hook_block).call(target_method, *args, **kwargs, &block)
+        end
       end
     end
 

sorbet/rbi/shims/arel.rbi

@@ -0,0 +1,24 @@
+# typed: strict
+# frozen_string_literal: true
+
+class Arel::Attributes::Attribute
+  sig { returns(Arel::Table) }
+  def relation
+  end
+
+  sig { returns(String) }
+  def name
+  end
+end
+
+module ActiveRecord::QueryMethods
+  sig { returns(ActiveRecord::Relation::WhereClause) }
+  def where_clause
+  end
+end
+
+class ActiveRecord::Relation::WhereClause
+  sig { returns(T::Array[Arel::Attributes::Attribute]) }
+  def extract_attributes
+  end
+end

spec/has_aux_table_spec.rb

@@ -234,13 +234,13 @@ RSpec.describe HasAuxTable do
 
     it "allows saving the model with auxiliary columns" do
       car = Car.create!(name: "Honda Civic")
-      num_queries =
+      queries =
-        SpecHelper.count_queries do
+        SpecHelper.capture_queries do
           car.fuel_type = "hybrid"
           car.engine_size = 1.8
           car.save!
         end
-      expect(num_queries).to eq(1)
+      expect(queries.length).to eq(1)
     end
   end
 
@@ -423,32 +423,32 @@ RSpec.describe HasAuxTable do
         # All query methods now use single queries with proper LEFT OUTER JOINs
 
         it "loads single model with auxiliary data in one query using find" do
-          query_count =
+          queries =
-            SpecHelper.count_queries do
+            SpecHelper.capture_queries do
               car = Car.find(@car1.id)
               # Access auxiliary attributes to ensure they're loaded
               car.fuel_type
               car.engine_size
             end
 
-          expect(query_count).to eq(1)
+          expect(queries.length).to eq(1)
         end
 
         it "loads single model with auxiliary data in one query using find_by" do
-          query_count =
+          queries =
-            SpecHelper.count_queries do
+            SpecHelper.capture_queries do
               car = Car.find_by(name: "Toyota Prius")
               # Access auxiliary attributes to ensure they're loaded
               car.fuel_type
               car.engine_size
             end
 
-          expect(query_count).to eq(1)
+          expect(queries.length).to eq(1)
         end
 
         it "loads multiple models with auxiliary data in one query using where" do
-          query_count =
+          queries =
-            SpecHelper.count_queries do
+            SpecHelper.capture_queries do
               cars = Car.where(fuel_type: %w[hybrid electric])
               # Access auxiliary attributes for all cars
               cars.each do |car|
@@ -457,7 +457,7 @@ RSpec.describe HasAuxTable do
               end
             end
 
-          expect(query_count).to eq(1)
+          expect(queries.length).to eq(1)
         end
 
         it "avoids N+1 queries when loading multiple models" do
@@ -472,8 +472,8 @@ RSpec.describe HasAuxTable do
           end
 
           cars = nil
-          query_count =
+          queries =
-            SpecHelper.count_queries do
+            SpecHelper.capture_queries do
               cars = Car.where(fuel_type: "gasoline")
               # Access auxiliary attributes for all cars - should not trigger additional queries
               cars.each do |car|
@@ -483,13 +483,13 @@ RSpec.describe HasAuxTable do
               end
             end
 
-          expect(query_count).to eq(1) # Single query regardless of how many cars are loaded
+          expect(queries.length).to eq(1) # Single query regardless of how many cars are loaded
           expect(cars.length).to be >= 1 # At least the original Honda Civic plus new cars
         end
 
         it "uses single query when ordering by auxiliary columns" do
-          query_count =
+          queries =
-            SpecHelper.count_queries do
+            SpecHelper.capture_queries do
               cars = Car.where(engine_size: 1.0..3.0).order(:engine_size)
               # Access all attributes
               cars.each do |car|
@@ -499,12 +499,12 @@ RSpec.describe HasAuxTable do
               end
             end
 
-          expect(query_count).to eq(1)
+          expect(queries.length).to eq(1)
         end
 
         it "uses single query for complex auxiliary column queries" do
-          query_count =
+          queries =
-            SpecHelper.count_queries do
+            SpecHelper.capture_queries do
               cars =
                 Car.where(fuel_type: "hybrid").or(Car.where(engine_size: 0.0))
               # Access all attributes
@@ -515,12 +515,12 @@ RSpec.describe HasAuxTable do
               end
             end
 
-          expect(query_count).to eq(1)
+          expect(queries.length).to eq(1)
         end
 
         it "uses single query when finding by auxiliary columns" do
-          query_count =
+          queries =
-            SpecHelper.count_queries do
+            SpecHelper.capture_queries do
               car = Car.find_by(fuel_type: "hybrid", name: "Toyota Prius")
               # Access all attributes
               car.name
@@ -528,7 +528,7 @@ RSpec.describe HasAuxTable do
               car.engine_size
             end
 
-          expect(query_count).to eq(1)
+          expect(queries.length).to eq(1)
         end
 
         it "doesn't trigger additional queries when accessing auxiliary attributes after load" do
@@ -536,8 +536,8 @@ RSpec.describe HasAuxTable do
           car = Car.find(@car1.id)
 
           # Now count queries when accessing auxiliary attributes
-          query_count =
+          queries =
-            SpecHelper.count_queries do
+            SpecHelper.capture_queries do
               car.fuel_type
               car.engine_size
               car.fuel_type? # presence check
@@ -545,12 +545,12 @@ RSpec.describe HasAuxTable do
             end
 
           # Currently this should be 0 since auxiliary record is already loaded
-          expect(query_count).to eq(0) # No additional queries should be triggered
+          expect(queries.length).to eq(0) # No additional queries should be triggered
         end
 
         it "handles mixed queries with main and auxiliary columns in single query" do
-          query_count =
+          queries =
-            SpecHelper.count_queries do
+            SpecHelper.capture_queries do
               cars = Car.where(name: "Toyota Prius", fuel_type: "hybrid")
               cars.each do |car|
                 car.name
@@ -559,12 +559,12 @@ RSpec.describe HasAuxTable do
               end
             end
 
-          expect(query_count).to eq(1)
+          expect(queries.length).to eq(1)
         end
 
         it "uses single query for range queries on auxiliary columns" do
-          query_count =
+          queries =
-            SpecHelper.count_queries do
+            SpecHelper.capture_queries do
               cars = Car.where(engine_size: 0.0..1.9)
               cars.each do |car|
                 car.name
@@ -573,19 +573,19 @@ RSpec.describe HasAuxTable do
               end
             end
 
-          expect(query_count).to eq(1)
+          expect(queries.length).to eq(1)
         end
 
         it "maintains single query performance with limit and offset" do
-          query_count =
+          queries =
-            SpecHelper.count_queries do
+            SpecHelper.capture_queries do
               car = Car.where(fuel_type: %w[hybrid electric]).limit(1).first
               # Access auxiliary attributes
               car.fuel_type
               car.engine_size
             end
 
-          expect(query_count).to eq(1)
+          expect(queries.length).to eq(1)
         end
       end
 
@@ -983,8 +983,8 @@ RSpec.describe HasAuxTable do
       end
 
       it "reloads with one query" do
-        num_queries = SpecHelper.count_queries { @car.reload }
+        queries = SpecHelper.capture_queries { @car.reload }
-        expect(num_queries).to eq(1)
+        expect(queries.length).to eq(1)
       end
 
       it "reloads associations" do
@@ -1027,12 +1027,16 @@ RSpec.describe HasAuxTable do
     expect(Car.count).to eq(1)
     expect(Boat.count).to eq(1)
 
-    expect(SpecHelper.count_queries { car = Vehicle.find(car.id) }).to eq(1)
+    expect(
+      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(SpecHelper.count_queries { boat = Vehicle.find(boat.id) }).to eq(1)
+    expect(
+      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

@@ -0,0 +1,36 @@
+# typed: false
+# frozen_string_literal: true
+
+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")
+    end
+
+    it "queries only the aux table if plucking values that are on the aux table" do
+      queries =
+        SpecHelper.capture_queries do
+          expect(Car.pluck(:fuel_type)).to eq(%w[gasoline hybrid electric])
+        end
+
+      expect(queries.length).to eq(1)
+      expect(queries.first).not_to include("JOIN")
+    end
+
+    it "queries both tables if main table column is referenced" do
+      queries =
+        SpecHelper.capture_queries do
+          rel = Car.where(name: "Toyota Camry")
+          rel = rel.pluck(:fuel_type)
+          expect(rel).to eq(%w[gasoline])
+        end
+
+      expect(queries.length).to eq(1)
+      expect(queries.first).to include("JOIN")
+    end
+  end
+end

spec/spec_helper.rb

@@ -62,11 +62,13 @@ module SpecHelper
   LOG_QUERIES = T.let(false, T::Boolean)
 
   # Helper method to count queries
-  sig { params(block: T.proc.void).returns(Integer) }
+  sig { params(block: T.proc.void).returns(T::Array[String]) }
-  def self.count_queries(&block)
+  def self.capture_queries(&block)
-    query_count = 0
+    queries = T.let([], T::Array[String])
     query_callback =
-      lambda { |name, start, finish, message_id, values| query_count += 1 }
+      lambda do |name, start, finish, message_id, values|
+        queries << values[:sql]
+      end
 
     ActiveSupport::Notifications.subscribed(
       query_callback,
@@ -81,6 +83,6 @@ module SpecHelper
       ActiveRecord::Base.logger = old_logger if LOG_QUERIES
     end
 
-    query_count
+    queries
   end
 end