1 files changed, 10 insertions, 5 deletions
diff --git a/lib/gitlab/database/background_migration/batched_migration.rb b/lib/gitlab/database/background_migration/batched_migration.rb
index a883996a5c5..83beee091f1 100644
--- a/lib/gitlab/database/background_migration/batched_migration.rb
+++ b/lib/gitlab/database/background_migration/batched_migration.rb
@@ -79,6 +79,10 @@ module Gitlab
             transition any => :finalizing
           end
 
+          before_transition any => :finished do |migration|
+            migration.finished_at = Time.current if migration.respond_to?(:finished_at)
+          end
+
           before_transition any => :active do |migration|
             migration.started_at = Time.current if migration.respond_to?(:started_at)
           end
@@ -92,10 +96,6 @@ module Gitlab
           for_configuration(gitlab_schema, job_class_name, table_name, column_name, job_arguments).first
         end
 
-        def self.active_migration(connection:)
-          active_migrations_distinct_on_table(connection: connection, limit: 1).first
-        end
-
         def self.find_executable(id, connection:)
           for_gitlab_schema(Gitlab::Database.gitlab_schemas_for_connection(connection))
             .executable.find_by_id(id)
@@ -220,7 +220,12 @@ module Gitlab
         end
 
         def health_context
-          HealthStatus::Context.new(connection, [table_name], gitlab_schema.to_sym)
+          @health_context ||= Gitlab::Database::HealthStatus::Context.new(
+            self,
+            connection,
+            [table_name],
+            gitlab_schema.to_sym
+          )
         end
 
         def hold!(until_time: 10.minutes.from_now)