Fixes issues with update_all, increment, and decrement

lib/activerecord-multi-tenant/relation_extension.rb

@@ -19,9 +19,21 @@ 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?
 
+      if updates.is_a?(Hash)
+        if klass.locking_enabled? &&
+           !updates.key?(klass.locking_column) &&
+           !updates.key?(klass.locking_column.to_sym)
+          attr = table[klass.locking_column]
+          updates[attr.name] = _increment_attribute(attr)
+        end
+        values = _substitute_values(updates)
+      else
+        values = Arel.sql(klass.sanitize_sql_for_assignment(updates, table.name))
+      end
+
       stmt = Arel::UpdateManager.new
       stmt.table(table)
-      stmt.set Arel.sql(klass.send(:sanitize_sql_for_assignment, updates))
+      stmt.set values
       stmt.wheres = [generate_in_condition_subquery]
 
       klass.connection.update(stmt, "#{klass} Update All").tap { reset }

spec/activerecord-multi-tenant/query_rewriter_spec.rb

@@ -63,6 +63,9 @@
       @queries.each do |actual_query|
         next unless actual_query.include?('UPDATE "projects" SET "name"')
 
+        # Extract parameterized values and replace placeholders
+        actual_query = actual_query.gsub('$1', "'New Name'")
+
         expect(format_sql(actual_query)).to eq(format_sql(expected_query.gsub(':account_id', account.id.to_s)))
       end
     end
@@ -99,6 +102,9 @@
       @queries.each do |actual_query|
         next unless actual_query.include?('UPDATE "projects" SET "name"')
 
+        # Extract parameterized values and replace placeholders
+        actual_query = actual_query.gsub('$1', "'#{new_name}'").gsub('$2', limit.to_s)
+
         expect(format_sql(actual_query.gsub('$1',
                                             limit.to_s)).strip).to eq(format_sql(expected_query).strip)
       end
@@ -245,4 +251,54 @@
       ).to eq([page_in_alive_domain.id])
     end
   end
+
+  context 'when updating locked column in comments' do
+    let!(:account) { Account.create!(name: 'Test Account') }
+    let!(:comment1) { Comment.create!(account_id: account.id, commentable_id: 1, commentable_type: 'Post') }
+    let!(:comment2) { Comment.create!(account_id: account.id, commentable_id: 2, commentable_type: 'Post') }
+
+    it 'updates the locked column for all records' do
+      expect do
+        Comment.update_all(locked: 1)
+      end.to change { Comment.where(locked: 1).count }.from(0).to(2)
+    end
+
+    it 'update_all the records with expected query' do
+      expected_query = <<-SQL.strip
+        UPDATE "comments" SET "locked" = 1 WHERE "comments"."id" IN
+          (SELECT "comments"."id" FROM "comments"
+              WHERE "comments"."account_id" = :account_id
+          )
+          AND "comments"."account_id" = :account_id
+      SQL
+
+      expect do
+        MultiTenant.with(account) do
+          Comment.update_all(locked: 1)
+        end
+      end.to change { Comment.where(locked: 1).count }.from(0).to(2)
+
+      @queries.each do |actual_query|
+        next unless actual_query.include?('UPDATE "comments" SET "locked"')
+
+        # Extract parameterized values and replace placeholders
+        actual_query = actual_query.gsub('$1', '1')
+
+        expect(format_sql(actual_query)).to eq(format_sql(expected_query.gsub(':account_id', account.id.to_s)))
+      end
+    end
+  end
+
+  context 'when using decrement or increment' do
+    let!(:account) { Account.create!(name: 'Test Account') }
+    let!(:comment) { Comment.create!(account_id: account.id, commentable_id: 1, commentable_type: 'Post', locked: 1) }
+
+    it 'increments and decrements the column `locked`' do
+      expect do
+        comment.increment!(:locked)
+        comment.increment!(:locked)
+        comment.decrement!(:locked)
+      end.to change { comment.reload.locked }.from(1).to(2)
+    end
+  end
 end

spec/schema.rb

@@ -78,6 +78,7 @@
     t.column :account_id, :integer
     t.column :commentable_id, :integer
     t.column :commentable_type, :string
+    t.column :locked, :integer, default: 0
   end
 
   create_table :partition_key_not_model_tasks, force: true, partition_key: :non_model_id do |t|