Add ensure_background_migration_succeeded migration helper for back…

…ground migrations

CHANGELOG.md

@@ -1,5 +1,7 @@
 ## master (unreleased)
 
+- Add `ensure_background_migration_succeeded` migration helper for background migrations
+
 ## 0.14.0 (2024-02-01)
 
 - Add ability to configure whether background migrations should be run inline

docs/background_migrations.md

@@ -125,6 +125,10 @@ enqueue_background_migration("MyMigrationWithArgs", arg1, arg2, ...)
 
 **Note**: These migration helpers should be run inside the migration against the database where background migrations tables are defined.
 
+## Depending on migrated data
+
+You shouldn't depend on the data until the background migration is finished. If having 100% of the data migrated is a requirement, then the `ensure_background_migration_succeeded` helper can be used to guarantee that the migration was succeeded and the data fully migrated.
+
 ## Testing
 
 At a minimum, it's recommended that the `#process_batch` method in your background migration is tested. You may also want to test the `#relation` and `#count` methods if they are sufficiently complex.

lib/online_migrations/background_migrations/migration.rb

@@ -20,6 +20,7 @@ class Migration < ApplicationRecord
       self.table_name = :background_migrations
 
       scope :queue_order, -> { order(created_at: :asc) }
+      scope :parents, -> { where(parent_id: nil) }
       scope :runnable, -> { where(composite: false) }
       scope :active, -> { where(status: [statuses[:enqueued], statuses[:running]]) }
       scope :except_succeeded, -> { where.not(status: :succeeded) }

lib/online_migrations/background_migrations/migration_helpers.rb

@@ -381,6 +381,42 @@ def enqueue_background_migration(migration_name, *arguments, **options)
         migration
       end
 
+      # Ensures that the background migration with provided configuration is succeeded.
+      #
+      # If the enqueued migration was not found (probably when resetting a dev environment
+      # followed by `db:migrate`) then a log warning is printed.
+      # If enqueued migration was found but is not succeeded, then the error is raised.
+      #
+      # @param migration_name [String, Class] Background migration job class name
+      # @param arguments [Array, nil] Arguments with which background migration was enqueued
+      #
+      # @example Without arguments
+      #   ensure_background_migration_succeeded("BackfillProjectIssuesCount")
+      #
+      # @example With arguments
+      #   ensure_background_migration_succeeded("CopyColumn", arguments: ["users", "id", "id_for_type_change"])
+      #
+      def ensure_background_migration_succeeded(migration_name, arguments: nil)
+        migration_name = migration_name.name if migration_name.is_a?(Class)
+
+        configuration = { migration_name: migration_name }
+
+        if arguments
+          arguments = Array(arguments)
+          migration = Migration.parents.for_configuration(migration_name, arguments).first
+          configuration[:arguments] = arguments.to_json
+        else
+          migration = Migration.parents.for_migration_name(migration_name).first
+        end
+
+        if migration.nil?
+          Utils.say("Could not find background migration for the given configuration: #{configuration}")
+        elsif !migration.succeeded?
+          raise "Expected background migration for the given configuration to be marked as 'succeeded', " \
+                "but it is '#{migration.status}': #{configuration}"
+        end
+      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,

lib/online_migrations/migration.rb

@@ -12,6 +12,7 @@ def migrate(direction)
       super
     ensure
       VerboseSqlLogs.disable if verbose_sql_logs?
+      OnlineMigrations.current_migration = nil
     end
 
     # @private

test/background_migrations/background_migrations.rb

@@ -24,6 +24,8 @@ class Dog < ShardRecord
   end
 
   class MakeAllNonAdmins < OnlineMigrations::BackgroundMigration
+    def initialize(*_dummy_args) end
+
     def relation
       User.all
     end

test/schema_statements/misc_test.rb

@@ -203,6 +203,42 @@ def test_run_background_migrations_inline_configured_to_custom_proc
       assert_nil user.reload.admin
     end
 
+    def test_ensure_background_migration_succeeded
+      m = @connection.enqueue_background_migration("MakeAllNonAdmins")
+      assert m.succeeded?
+      assert_nothing_raised do
+        @connection.ensure_background_migration_succeeded("MakeAllNonAdmins")
+      end
+    end
+
+    def test_ensure_background_migration_succeeded_with_arguments
+      m = @connection.enqueue_background_migration("MakeAllNonAdmins", 1, { foo: "bar" }, 2)
+      assert m.succeeded?
+      assert_nothing_raised do
+        @connection.ensure_background_migration_succeeded("MakeAllNonAdmins", arguments: [1, { foo: "bar" }, 2])
+      end
+    end
+
+    def test_ensure_background_migration_succeeded_when_migration_not_found
+      out, = capture_io do
+        logger = ActiveSupport::Logger.new($stdout)
+        ActiveRecord::Base.stub(:logger, logger) do
+          @connection.ensure_background_migration_succeeded("MakeAllNonAdmins")
+        end
+      end
+      assert_match(/Could not find background migration/i, out)
+    end
+
+    def test_ensure_background_migration_succeeded_when_migration_is_failed
+      m = @connection.enqueue_background_migration("MakeAllNonAdmins")
+      m.update_column(:status, :failed)
+
+      error = assert_raises(RuntimeError) do
+        @connection.ensure_background_migration_succeeded("MakeAllNonAdmins")
+      end
+      assert_match(/to be marked as 'succeeded'/i, error.message)
+    end
+
     def test_disable_statement_timeout
       prev_value = get_statement_timeout
       set_statement_timeout(10)