From e3d2c88140e635fc9dbb75e87378d80484d419f3 Mon Sep 17 00:00:00 2001 From: amit909singh Date: Thu, 4 Jan 2024 23:23:41 -0800 Subject: [PATCH] Fix the tenant scoping for update_all (#223) --- ...te_operations.rb => relation_extension.rb} | 40 +++++-- lib/activerecord_multi_tenant.rb | 2 +- .../query_rewriter_spec.rb | 102 ++++++++++++++++++ 3 files changed, 136 insertions(+), 8 deletions(-) rename lib/activerecord-multi-tenant/{delete_operations.rb => relation_extension.rb} (50%) diff --git a/lib/activerecord-multi-tenant/delete_operations.rb b/lib/activerecord-multi-tenant/relation_extension.rb similarity index 50% rename from lib/activerecord-multi-tenant/delete_operations.rb rename to lib/activerecord-multi-tenant/relation_extension.rb index b2bfbf3..c1d8900 100644 --- a/lib/activerecord-multi-tenant/delete_operations.rb +++ b/lib/activerecord-multi-tenant/relation_extension.rb @@ -7,29 +7,55 @@ 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 = eager_loading? ? apply_join_dependency.arel : build_arel 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) - # Check if the tenant key is present in the model's column names tenant_condition = table[tenant_key].eq(tenant_id) - # Add the tenant condition to the arel query if it is not already present unless arel.constraints.any? { |node| node.to_sql.include?(tenant_condition.to_sql) } arel = arel.where(tenant_condition) 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]) - in_condition = Arel::Nodes::In.new(table[primary_key], subquery.ast) - stmt = Arel::DeleteManager.new.from(table) - stmt.wheres = [in_condition] - # Execute the delete statement using the connection and return the result - klass.connection.delete(stmt, "#{klass} Delete All").tap { reset } + # Create an IN condition node with the primary key of the table and the subquery + Arel::Nodes::In.new(table[primary_key], subquery.ast) end end end diff --git a/lib/activerecord_multi_tenant.rb b/lib/activerecord_multi_tenant.rb index 09a29a0..7c4fc30 100644 --- a/lib/activerecord_multi_tenant.rb +++ b/lib/activerecord_multi_tenant.rb @@ -11,4 +11,4 @@ require_relative 'activerecord-multi-tenant/table_node' require_relative 'activerecord-multi-tenant/version' require_relative 'activerecord-multi-tenant/habtm' -require_relative 'activerecord-multi-tenant/delete_operations' +require_relative 'activerecord-multi-tenant/relation_extension' diff --git a/spec/activerecord-multi-tenant/query_rewriter_spec.rb b/spec/activerecord-multi-tenant/query_rewriter_spec.rb index 2b6410e..db0edda 100644 --- a/spec/activerecord-multi-tenant/query_rewriter_spec.rb +++ b/spec/activerecord-multi-tenant/query_rewriter_spec.rb @@ -3,6 +3,13 @@ require 'spec_helper' describe 'Query Rewriter' do + before(:each) do + @queries = [] + ActiveSupport::Notifications.subscribe('sql.active_record') do |_name, _started, _finished, _unique_id, payload| + @queries << payload[:sql] + end + end + context 'when bulk updating' do let!(:account) { Account.create!(name: 'Test Account') } let!(:project) { Project.create(name: 'Project 1', account: account) } @@ -35,6 +42,67 @@ project.update(name: 'New Name') end.to change { project.reload.name }.from('Project 1').to('New Name') end + + it 'update_all the records with expected query' do + expected_query = <<-SQL.strip + 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 + WHERE "projects"."account_id" = :account_id + ) + AND "projects"."account_id" = :account_id + SQL + + expect do + MultiTenant.with(account) do + Project.joins(:manager).update_all(name: 'New Name') + end + end.to change { project.reload.name }.from('Project 1').to('New Name') + + @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))) + end + end + + it 'updates a limited number of records with expected query' do + # create 2 more projects + Project.create(name: 'project2', account: account) + Project.create(name: 'project3', account: account) + new_name = 'New Name' + limit = 2 + expected_query = <<-SQL + UPDATE + "projects" + SET + "name" = 'New Name' + WHERE + "projects"."id" IN ( + SELECT + "projects"."id" + FROM + "projects" + WHERE + "projects"."account_id" = #{account.id} LIMIT #{limit} + ) + AND "projects"."account_id" = #{account.id} + SQL + + expect do + MultiTenant.with(account) do + Project.limit(limit).update_all(name: new_name) + end + end.to change { Project.where(name: new_name).count }.from(0).to(limit) + + @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) + end + end end context 'when bulk deleting' do @@ -102,6 +170,40 @@ end.to change { Project.count }.from(3).to(1) end + it 'deletes a limited number of records with expected query' do + # create 2 more projects + Project.create(name: 'project2', account: account) + Project.create(name: 'project3', account: account) + limit = 2 + expected_query = <<-SQL + DELETE FROM + "projects" + WHERE + "projects"."id" IN ( + SELECT + "projects"."id" + FROM + "projects" + WHERE + "projects"."account_id" = #{account.id} LIMIT #{limit} + ) + AND "projects"."account_id" = #{account.id} + SQL + + expect do + MultiTenant.with(account) do + Project.limit(limit).delete_all + end + end.to change { Project.count }.by(-limit) + + @queries.each do |actual_query| + next unless actual_query.include?('DELETE FROM "projects"') + + expect(format_sql(actual_query.gsub('$1', + limit.to_s)).strip).to eq(format_sql(expected_query).strip) + end + end + it 'destroy the record' do expect do MultiTenant.with(account) do