1 files changed, 17 insertions, 9 deletions
diff --git a/lib/gitlab/database/background_migration/batched_migration_wrapper.rb b/lib/gitlab/database/background_migration/batched_migration_wrapper.rb
index e37df102872..057f856d859 100644
--- a/lib/gitlab/database/background_migration/batched_migration_wrapper.rb
+++ b/lib/gitlab/database/background_migration/batched_migration_wrapper.rb
@@ -6,6 +6,10 @@ module Gitlab
       class BatchedMigrationWrapper
         extend Gitlab::Utils::StrongMemoize
 
+        def initialize(connection: ApplicationRecord.connection)
+          @connection = connection
+        end
+
         # Wraps the execution of a batched_background_migration.
         #
         # Updates the job's tracking records with the status of the migration
@@ -18,24 +22,25 @@ module Gitlab
 
           execute_batch(batch_tracking_record)
 
-          batch_tracking_record.status = :succeeded
-        rescue Exception # rubocop:disable Lint/RescueException
-          batch_tracking_record.status = :failed
+          batch_tracking_record.succeed!
+        rescue Exception => error # rubocop:disable Lint/RescueException
+          batch_tracking_record.failure!(error: error)
 
           raise
         ensure
-          finish_tracking_execution(batch_tracking_record)
           track_prometheus_metrics(batch_tracking_record)
         end
 
         private
 
+        attr_reader :connection
+
         def start_tracking_execution(tracking_record)
-          tracking_record.update!(attempts: tracking_record.attempts + 1, status: :running, started_at: Time.current, finished_at: nil, metrics: {})
+          tracking_record.run!
         end
 
         def execute_batch(tracking_record)
-          job_instance = tracking_record.migration_job_class.new
+          job_instance = migration_instance_for(tracking_record.migration_job_class)
 
           job_instance.perform(
             tracking_record.min_value,
@@ -51,9 +56,12 @@ module Gitlab
           end
         end
 
-        def finish_tracking_execution(tracking_record)
-          tracking_record.finished_at = Time.current
-          tracking_record.save!
+        def migration_instance_for(job_class)
+          if job_class < Gitlab::BackgroundMigration::BaseJob
+            job_class.new(connection: connection)
+          else
+            job_class.new
+          end
         end
 
         def track_prometheus_metrics(tracking_record)