From 89fc190381df23ca116fce1441842901f6559055 Mon Sep 17 00:00:00 2001
From: fatkodima <fatkodima123@gmail.com>
Date: Tue, 7 May 2024 19:02:35 +0300
Subject: [PATCH] Fix setting `started_at`/`finished_at` for parents of sharded
 background schema migrations

---
 CHANGELOG.md                                        |  1 +
 .../background_schema_migrations/migration.rb       |  5 ++++-
 .../migration_runner.rb                             | 13 ++++++++++---
 .../migration_runner_test.rb                        |  5 +++++
 test/background_schema_migrations/migration_test.rb |  4 ++++
 5 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7fbe3ff..6a55c43 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,6 @@
 ## master (unreleased)
 
+- Fix setting `started_at`/`finished_at` for parents of sharded background schema migrations
 - Improve retrying of failed sharded background migrations
 - Fix a bug when retried background data migration can not start
 - Do not run multiple background schema migrations on the same table at the same time
diff --git a/lib/online_migrations/background_schema_migrations/migration.rb b/lib/online_migrations/background_schema_migrations/migration.rb
index 511dd11..8a782f1 100644
--- a/lib/online_migrations/background_schema_migrations/migration.rb
+++ b/lib/online_migrations/background_schema_migrations/migration.rb
@@ -111,7 +111,10 @@ def progress
       def retry
         if composite? && failed?
           children.failed.each(&:retry)
-          running!
+          update!(
+            status: self.class.statuses[:running],
+            finished_at: nil
+          )
           true
         elsif failed?
           update!(
diff --git a/lib/online_migrations/background_schema_migrations/migration_runner.rb b/lib/online_migrations/background_schema_migrations/migration_runner.rb
index faf6125..d123609 100644
--- a/lib/online_migrations/background_schema_migrations/migration_runner.rb
+++ b/lib/online_migrations/background_schema_migrations/migration_runner.rb
@@ -27,7 +27,14 @@ def run
         def mark_as_running
           Migration.transaction do
             migration.running!
-            migration.parent.running! if migration.parent
+
+            if (parent = migration.parent)
+              if parent.started_at
+                parent.update!(status: :running, finished_at: nil)
+              else
+                parent.update!(status: :running, started_at: Time.current, finished_at: nil)
+              end
+            end
           end
         end
 
@@ -90,10 +97,10 @@ def complete_parent_if_needed(migration)
           parent.with_lock do
             children = parent.children.select(:status)
             if children.all?(&:succeeded?)
-              parent.succeeded!
+              parent.update!(status: :succeeded, finished_at: Time.current)
               completed = true
             elsif children.any?(&:failed?)
-              parent.failed!
+              parent.update!(status: :failed, finished_at: Time.current)
               completed = true
             end
           end
diff --git a/test/background_schema_migrations/migration_runner_test.rb b/test/background_schema_migrations/migration_runner_test.rb
index bacec1a..0d3ffbd 100644
--- a/test/background_schema_migrations/migration_runner_test.rb
+++ b/test/background_schema_migrations/migration_runner_test.rb
@@ -27,7 +27,12 @@ def test_run_marks_migration_and_its_parent_as_running
 
       run_migration(child)
       assert child.reload.succeeded?
+      assert child.started_at
+      assert child.finished_at
+
       assert m.reload.running?
+      assert m.started_at
+      assert_nil m.finished_at
     end
 
     def test_run_runs_migration
diff --git a/test/background_schema_migrations/migration_test.rb b/test/background_schema_migrations/migration_test.rb
index a65afd7..522bdc7 100644
--- a/test/background_schema_migrations/migration_test.rb
+++ b/test/background_schema_migrations/migration_test.rb
@@ -175,6 +175,8 @@ def test_retry_composite
       m = create_sharded_migration
       run_migration(m)
       assert m.succeeded?
+      assert m.started_at
+      assert m.finished_at
 
       child1, child2, child3 = m.children.to_a
       migrations = [m, child1, child2, child3]
@@ -185,6 +187,8 @@ def test_retry_composite
       child1.update_column(:status, "failed")
 
       assert m.retry
+      assert m.started_at
+      assert_nil m.finished_at
 
       migrations.each(&:reload)
       assert m.running?