diff --git a/lib/has_aux_table/model_class_helper.rb b/lib/has_aux_table/model_class_helper.rb index 711eeef..4335562 100644 --- a/lib/has_aux_table/model_class_helper.rb +++ b/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 ||= diff --git a/lib/has_aux_table/relation_extensions.rb b/lib/has_aux_table/relation_extensions.rb index c7b08cc..89a1d87 100644 --- a/lib/has_aux_table/relation_extensions.rb +++ b/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] diff --git a/lib/has_aux_table/util.rb b/lib/has_aux_table/util.rb index d60598b..6767d93 100644 --- a/lib/has_aux_table/util.rb +++ b/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 - T.unsafe(hook_block).call(method, *args, **kwargs, &block) + if is_instance_method + 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 diff --git a/sorbet/rbi/shims/arel.rbi b/sorbet/rbi/shims/arel.rbi new file mode 100644 index 0000000..68e0125 --- /dev/null +++ b/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 diff --git a/spec/has_aux_table_spec.rb b/spec/has_aux_table_spec.rb index 18f98b9..b60e6cb 100644 --- a/spec/has_aux_table_spec.rb +++ b/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 = - SpecHelper.count_queries do + queries = + 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 = - SpecHelper.count_queries do + queries = + 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 = - SpecHelper.count_queries do + queries = + 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 = - SpecHelper.count_queries do + queries = + 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 = - SpecHelper.count_queries do + queries = + 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 = - SpecHelper.count_queries do + queries = + 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 = - SpecHelper.count_queries do + queries = + 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 = - SpecHelper.count_queries do + queries = + 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 = - SpecHelper.count_queries do + queries = + 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 = - SpecHelper.count_queries do + queries = + 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 = - SpecHelper.count_queries do + queries = + 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 = - SpecHelper.count_queries do + queries = + 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 } - expect(num_queries).to eq(1) + queries = SpecHelper.capture_queries { @car.reload } + 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 diff --git a/spec/loading_optimizations_spec.rb b/spec/loading_optimizations_spec.rb new file mode 100644 index 0000000..6564a8a --- /dev/null +++ b/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 diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb index bd414dd..de757c7 100644 --- a/spec/spec_helper.rb +++ b/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) } - def self.count_queries(&block) - query_count = 0 + sig { params(block: T.proc.void).returns(T::Array[String]) } + def self.capture_queries(&block) + 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