Fix the tenant scoping for update_all (#223)

Amit909Singh · web-flow · commit e3d2c88140e6 · 2024-01-05T10:23:41.000+03:00
diff --git a/lib/activerecord-multi-tenant/relation_extension.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
@@ -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
@@ -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
