diff options
Diffstat (limited to 'lib/gitlab/database/background_migration/batched_migration_wrapper.rb')
-rw-r--r-- | lib/gitlab/database/background_migration/batched_migration_wrapper.rb | 26 |
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) |