Make applying partition key changes more robust

lib/sequent/internal/partition_key_change.rb

@@ -12,15 +12,20 @@ class PartitionKeyChange < Sequent::ApplicationRecord
       belongs_to :partitioned_aggregate, primary_key: :aggregate_id, foreign_key: :aggregate_id
 
       def self.update_aggregate_partition_keys(limit:)
+        count = 0
         limit.times do
           ActiveRecord::Base.transaction do
             if Sequent::Migrations::Versions.running.present?
               fail Sequent::Migrations::ConcurrentMigration,
                    'cannot update partition keys while view schema migration is running'
             end
 
-            change = PartitionKeyChange.first
-            return unless change
+            PartitionKeyChange.connection.exec_update("SET LOCAL statement_timeout TO '30s'", 'statement_timeout')
+
+            change = PartitionKeyChange.lock('FOR UPDATE SKIP LOCKED').first
+            return count unless change
+
+            count += 1
 
             PartitionedAggregate
               .where(aggregate_id: change.aggregate_id)
@@ -30,6 +35,7 @@ def self.update_aggregate_partition_keys(limit:)
             change.destroy!
           end
         end
+        count
       end
     end
   end

lib/sequent/rake/migration_tasks.rb

@@ -308,11 +308,12 @@ def register_tasks!
                      'usage rake sequent:aggregates:update_partition_keys[limit]'
               end
 
-              count = Sequent::Internal::PartitionKeyChange.count
-              Sequent.logger.info "Applying #{count} partition key changes (limited to #{limit})"
+              total_count = Sequent::Internal::PartitionKeyChange.count
+              Sequent.logger.info "Applying #{total_count} partition key changes (limited to #{limit})"
 
               begin
-                Sequent::Internal::PartitionKeyChange.update_aggregate_partition_keys(limit:)
+                applied_count = Sequent::Internal::PartitionKeyChange.update_aggregate_partition_keys(limit:)
+                Sequent.logger.info "Applied #{applied_count} out of #{total_count} partition key changes"
               rescue Sequent::Migrations::ConcurrentMigration
                 Sequent.logger.error 'View schema migration is active so not updating partition keys'
               end

spec/lib/sequent/internal/partitioned_storage_spec.rb

@@ -80,7 +80,8 @@ module Internal
         end
 
         it 'updates the aggregate and events using a maintenance task' do
-          PartitionKeyChange.update_aggregate_partition_keys(limit: 10)
+          applied_count = PartitionKeyChange.update_aggregate_partition_keys(limit: 10)
+          expect(applied_count).to eq(1)
 
           logged_change = PartitionKeyChange.find_by(aggregate_id:)
           expect(logged_change).to be_nil