Fix a bug that can lead to unfinished children of a sharded backgroun…

…d migration

CHANGELOG.md

@@ -1,5 +1,7 @@
 ## master (unreleased)
 
+- Fix a bug that can lead to unfinished children of a sharded background migration
+
 ## 0.11.1 (2024-01-11)
 
 - Fix calculation of batch ranges for sharded background migrations

lib/online_migrations/background_migrations/migration.rb

@@ -2,6 +2,11 @@
 
 module OnlineMigrations
   module BackgroundMigrations
+    # Class representing background data migration.
+    #
+    # @note The records of this class should not be created manually, but via
+    #   `enqueue_background_migration` helper inside migrations.
+    #
     class Migration < ApplicationRecord
       STATUSES = [
         :enqueued,    # The migration has been enqueued by the user.
@@ -47,7 +52,6 @@ class Migration < ApplicationRecord
       validates_with MigrationStatusValidator, on: :update
 
       before_validation :set_defaults
-      before_create :create_child_migrations, if: :composite?
       before_update :copy_attributes_to_children, if: :composite?
 
       # @private
@@ -57,6 +61,7 @@ def self.normalize_migration_name(migration_name)
       end
 
       def migration_name=(class_name)
+        class_name = class_name.name if class_name.is_a?(Class)
         write_attribute(:migration_name, self.class.normalize_migration_name(class_name))
       end
 
@@ -209,10 +214,6 @@ def next_batch_range
         [min_value, max_value]
       end
 
-      protected
-        attr_accessor :child
-        alias child? child
-
       private
         def validate_batch_column_values
           if max_value.to_i < min_value.to_i
@@ -242,11 +243,6 @@ def validate_jobs_status
 
         def set_defaults
           if migration_relation.is_a?(ActiveRecord::Relation)
-            if !child?
-              shards = Utils.shard_names(migration_model)
-              self.composite = shards.size > 1
-            end
-
             self.batch_column_name ||= migration_relation.primary_key
 
             if composite?
@@ -276,18 +272,6 @@ def set_defaults
           self.batch_max_attempts   ||= config.batch_max_attempts
         end
 
-        def create_child_migrations
-          shards = Utils.shard_names(migration_model)
-
-          children = shards.map do |shard|
-            child = Migration.new(migration_name: migration_name, arguments: arguments, shard: shard)
-            child.child = true
-            child
-          end
-
-          self.children = children
-        end
-
         def copy_attributes_to_children
           attributes = [:batch_size, :sub_batch_size, :batch_pause, :sub_batch_pause_ms, :batch_max_attempts]
           updates = {}

lib/online_migrations/background_migrations/migration_helpers.rb

@@ -358,23 +358,41 @@ def perform_action_on_relation_in_background(model_name, conditions, action, upd
       #     in development and test environments
       #
       def enqueue_background_migration(migration_name, *arguments, **options)
+        migration = create_background_migration(migration_name, *arguments, **options)
+
+        # For convenience in dev/test environments
+        if Utils.developer_env?
+          runner = MigrationRunner.new(migration)
+          runner.run_all_migration_jobs
+        end
+
+        migration
+      end
+
+      # @private
+      def create_background_migration(migration_name, *arguments, **options)
         options.assert_valid_keys(:batch_column_name, :min_value, :max_value, :batch_size, :sub_batch_size,
             :batch_pause, :sub_batch_pause_ms, :batch_max_attempts)
 
-        migration_name = migration_name.name if migration_name.is_a?(Class)
-
-        migration = Migration.create!(
+        migration = Migration.new(
           migration_name: migration_name,
           arguments: arguments,
           **options
         )
 
-        # For convenience in dev/test environments
-        if Utils.developer_env?
-          runner = MigrationRunner.new(migration)
-          runner.run_all_migration_jobs
+        shards = Utils.shard_names(migration.migration_model)
+        if shards.size > 1
+          migration.children = shards.map do |shard|
+            child = migration.dup
+            child.shard = shard
+            child
+          end
+
+          migration.composite = true
         end
 
+        # This will save all the records using a transaction.
+        migration.save!
         migration
       end
     end

test/background_migrations/migration_job_test.rb

@@ -70,10 +70,8 @@ def test_retry
     end
 
     private
-      def create_migration(attributes = {})
-        OnlineMigrations::BackgroundMigrations::Migration.create!(
-          { migration_name: "EachBatchCalled" }.merge(attributes)
-        )
+      def create_migration(migration_name: "EachBatchCalled", **attributes)
+        @connection.create_background_migration(migration_name, **attributes)
       end
 
       def run_migration_job(migration)

test/background_migrations/migration_runner_test.rb

@@ -224,6 +224,8 @@ def test_run_all_migration_jobs_on_composite_migration
       on_each_shard { 2.times { Dog.create! } }
 
       m = create_migration(migration_name: "MakeAllDogsNice", batch_size: 1, sub_batch_size: 1)
+      # Reload children. There was a bug, not shown in tests, because the association was cached.
+      m.children.reload
       run_all_migration_jobs(m)
 
       on_each_shard do
@@ -278,10 +280,8 @@ def test_finish_composite_migration
     end
 
     private
-      def create_migration(attributes = {})
-        OnlineMigrations::BackgroundMigrations::Migration.create!(
-          { migration_name: "MakeAllNonAdmins" }.merge(attributes)
-        )
+      def create_migration(migration_name: "MakeAllNonAdmins", **attributes)
+        @connection.create_background_migration(migration_name, **attributes)
       end
 
       def run_migration_job(migration)

test/background_migrations/migration_test.rb

@@ -78,7 +78,7 @@ def test_migration_relation_with_order_clause
     end
 
     def test_status_transitions
-      m = create_migration(status: :enqueued)
+      m = create_migration
 
       m.status = :succeeded
       assert_not m.valid?
@@ -149,12 +149,12 @@ def test_empty_relation
     end
 
     def test_progress_succeded_migration
-      m = create_migration(status: :succeeded)
+      m = build_migration(status: :succeeded)
       assert_in_delta 100.0, m.progress
     end
 
     def test_progress_succeded_sharded_migration
-      m = create_migration(migration_name: "MakeAllDogsNice", status: :succeeded)
+      m = build_migration(migration_name: "MakeAllDogsNice", status: :succeeded)
       assert_in_delta 100.0, m.progress
     end
 
@@ -383,10 +383,8 @@ def test_copies_attribute_changes_to_child_migrations
     end
 
     private
-      def create_migration(attributes = {})
-        migration = build_migration(attributes)
-        migration.save!
-        migration
+      def create_migration(migration_name: "MakeAllNonAdmins", **attributes)
+        @connection.create_background_migration(migration_name, **attributes)
       end
 
       def build_migration(attributes = {})