From 8f100c9a6299f51592d1eb938994041592f0fd57 Mon Sep 17 00:00:00 2001 From: a5-stable Date: Thu, 19 Feb 2026 00:48:13 +0900 Subject: [PATCH] Refactor tenant enforcement to use Arel visitor pattern --- .../query_rewriter.rb | 18 ++++- .../relation_extension.rb | 79 ++++--------------- .../query_rewriter_spec.rb | 17 ++-- 3 files changed, 43 insertions(+), 71 deletions(-) diff --git a/lib/activerecord-multi-tenant/query_rewriter.rb b/lib/activerecord-multi-tenant/query_rewriter.rb index e7e03c65..bd3b30a1 100644 --- a/lib/activerecord-multi-tenant/query_rewriter.rb +++ b/lib/activerecord-multi-tenant/query_rewriter.rb @@ -263,7 +263,10 @@ def join_to_delete(delete, *args) def update(arel, name = nil, binds = []) model = MultiTenant.multi_tenant_model_for_arel(arel) - if model.present? && !MultiTenant.with_write_only_mode_enabled? && MultiTenant.current_tenant_id.present? + if model.present? && + !MultiTenant.with_write_only_mode_enabled? && + MultiTenant.current_tenant_id.present? && + !already_has_tenant_enforcement_clause?(arel) arel.where(MultiTenant::TenantEnforcementClause.new(model.arel_table[model.partition_key])) end super @@ -271,11 +274,22 @@ def update(arel, name = nil, binds = []) def delete(arel, name = nil, binds = []) model = MultiTenant.multi_tenant_model_for_arel(arel) - if model.present? && !MultiTenant.with_write_only_mode_enabled? && MultiTenant.current_tenant_id.present? + if model.present? && + !MultiTenant.with_write_only_mode_enabled? && + MultiTenant.current_tenant_id.present? && + !already_has_tenant_enforcement_clause?(arel) arel.where(MultiTenant::TenantEnforcementClause.new(model.arel_table[model.partition_key])) end super end + + private + + def already_has_tenant_enforcement_clause?(arel) + arel.try(:ast).try(:wheres).to_a.any? do |where| + where.is_a?(MultiTenant::BaseTenantEnforcementClause) + end + end end end diff --git a/lib/activerecord-multi-tenant/relation_extension.rb b/lib/activerecord-multi-tenant/relation_extension.rb index 2e3d5e3c..860d0893 100644 --- a/lib/activerecord-multi-tenant/relation_extension.rb +++ b/lib/activerecord-multi-tenant/relation_extension.rb @@ -1,71 +1,26 @@ # frozen_string_literal: true -module Arel - module ActiveRecordRelationExtension - # Overrides the delete_all method to include tenant scoping - def delete_all - # Call the original delete_all method if the current tenant is identified by an ID - return super if MultiTenant.current_tenant_is_id? || MultiTenant.current_tenant.nil? - - stmt = Arel::DeleteManager.new.from(table) - stmt.wheres = [generate_in_condition_subquery] - - # Execute the delete statement using the connection and return the result - klass.connection.delete(stmt, "#{klass} Delete All").tap { reset } - end - - # Overrides the update_all method to include tenant scoping - def update_all(updates) - # Call the original update_all method if the current tenant is identified by an ID - return super if MultiTenant.current_tenant_is_id? || MultiTenant.current_tenant.nil? - - stmt = Arel::UpdateManager.new - stmt.table(table) - stmt.set Arel.sql(klass.send(:sanitize_sql_for_assignment, updates)) - stmt.wheres = [generate_in_condition_subquery] - - klass.connection.update(stmt, "#{klass} Update All").tap { reset } - end - - private - - # The generate_in_condition_subquery method generates a subquery that selects - # records associated with the current tenant. - def generate_in_condition_subquery - # Get the tenant key and tenant ID based on the current tenant - tenant_key = MultiTenant.partition_key(MultiTenant.current_tenant_class) - tenant_id = MultiTenant.current_tenant_id - - # Build an Arel query - arel = if eager_loading? - apply_join_dependency.arel - elsif ActiveRecord.gem_version >= Gem::Version.create('7.2.0') - build_arel(klass.connection) - else - build_arel - end - - arel.source.left = table - - # If the tenant ID is present and the tenant key is a column in the model, - # add a condition to only include records where the tenant key equals the tenant ID - if tenant_id && klass.column_names.include?(tenant_key) - tenant_condition = table[tenant_key].eq(tenant_id) - unless arel.constraints.any? { |node| node.to_sql.include?(tenant_condition.to_sql) } - arel = arel.where(tenant_condition) +module Arel # :nodoc: all + module Visitors + module ToSqlPatch + def prepare_update_statement(object) + if object.key && (has_limit_or_offset_or_orders?(object) || has_join_sources?(object)) + stmt = super + + model = MultiTenant.multi_tenant_model_for_table(MultiTenant::TableNode.table_name(object.relation.left)) + if model.present? && !MultiTenant.with_write_only_mode_enabled? && MultiTenant.current_tenant_id.present? + stmt.wheres << MultiTenant::TenantEnforcementClause.new(model.arel_table[model.partition_key]) + end + + stmt + else + super end end - # Clone the query, clear its projections, and set its projection to the primary key of the table - subquery = arel.clone - subquery.projections.clear - subquery = subquery.project(table[primary_key]) - - # Create an IN condition node with the primary key of the table and the subquery - Arel::Nodes::In.new(table[primary_key], subquery.ast) + alias prepare_delete_statement prepare_update_statement end end end -# Patch ActiveRecord::Relation with the extension module -ActiveRecord::Relation.prepend(Arel::ActiveRecordRelationExtension) +Arel::Visitors::ToSql.prepend(Arel::Visitors::ToSqlPatch) diff --git a/spec/activerecord-multi-tenant/query_rewriter_spec.rb b/spec/activerecord-multi-tenant/query_rewriter_spec.rb index db0edda8..27254077 100644 --- a/spec/activerecord-multi-tenant/query_rewriter_spec.rb +++ b/spec/activerecord-multi-tenant/query_rewriter_spec.rb @@ -45,7 +45,7 @@ it 'update_all the records with expected query' do expected_query = <<-SQL.strip - UPDATE "projects" SET "name" = 'New Name' WHERE "projects"."id" IN + UPDATE "projects" SET "name" = 'New Name' WHERE ("projects"."id") IN (SELECT "projects"."id" FROM "projects" INNER JOIN "managers" ON "managers"."project_id" = "projects"."id" and "managers"."account_id" = :account_id @@ -63,7 +63,9 @@ @queries.each do |actual_query| next unless actual_query.include?('UPDATE "projects" SET "name"') - expect(format_sql(actual_query)).to eq(format_sql(expected_query.gsub(':account_id', account.id.to_s))) + expect(format_sql(actual_query.gsub('$1', "'New Name'"))).to eq(format_sql(expected_query.gsub( + ':account_id', account.id.to_s + ))) end end @@ -79,7 +81,7 @@ SET "name" = 'New Name' WHERE - "projects"."id" IN ( + ("projects"."id") IN ( SELECT "projects"."id" FROM @@ -99,8 +101,9 @@ @queries.each do |actual_query| next unless actual_query.include?('UPDATE "projects" SET "name"') - expect(format_sql(actual_query.gsub('$1', - limit.to_s)).strip).to eq(format_sql(expected_query).strip) + expect( + format_sql(actual_query.gsub('$1', "'#{new_name}'").gsub('$2', limit.to_s)).strip + ).to eq(format_sql(expected_query).strip) end end end @@ -126,7 +129,7 @@ it 'delete_all the records' do expected_query = <<-SQL.strip - DELETE FROM "projects" WHERE "projects"."id" IN + DELETE FROM "projects" WHERE ("projects"."id") IN (SELECT "projects"."id" FROM "projects" INNER JOIN "managers" ON "managers"."project_id" = "projects"."id" and "managers"."account_id" = :account_id @@ -179,7 +182,7 @@ DELETE FROM "projects" WHERE - "projects"."id" IN ( + ("projects"."id") IN ( SELECT "projects"."id" FROM