Add ensure_background_data_migration_succeeded and `ensure_backgrou…

…nd_schema_migration_succeeded` migration helpers

CHANGELOG.md

@@ -1,5 +1,6 @@
 ## master (unreleased)
 
+- Add `ensure_background_data_migration_succeeded` and `ensure_background_schema_migration_succeeded` migration helpers
 - Raise in development when background index creation/removal was not enqueued
 - Suggest two migrations for adding foreign keys
 - Reraise errors when running background schema migrations inline

docs/background_data_migrations.md

@@ -140,6 +140,10 @@ end
 
 **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 data migration is finished. If having 100% of the data migrated is a requirement, then the `ensure_background_data_migration_succeeded` helper can be used to guarantee that the migration 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.

docs/background_schema_migrations.md

@@ -62,6 +62,10 @@ end
 
 `add_index_in_background`/`remove_index_in_background` accept additional configuration options which controls how the background schema migration is run. Check the [source code](https://github.com/fatkodima/online_migrations/blob/master/lib/online_migrations/background_schema_migrations/migration_helpers.rb) for the list of all available configuration options.
 
+## Depending on schema changes
+
+You shouldn't depend on the schema until the background schema migration is finished. If having the schema migrated is a requirement, then the `ensure_background_schema_migration_succeeded` helper can be used to guarantee that the migration succeeded and the schema change applied.
+
 ## Instrumentation
 
 Background schema migrations use the [ActiveSupport::Notifications](http://api.rubyonrails.org/classes/ActiveSupport/Notifications.html) API.

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

@@ -395,6 +395,44 @@ def remove_background_data_migration(migration_name, *arguments)
       end
       alias remove_background_migration remove_background_data_migration
 
+      # Ensures that the background data migration with the provided configuration succeeded.
+      #
+      # If the enqueued migration was not found in development (probably when resetting a dev environment
+      # followed by `db:migrate`), then a log warning is printed.
+      # If enqueued migration was not found in production, then the error is raised.
+      # 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_data_migration_succeeded("BackfillProjectIssuesCount")
+      #
+      # @example With arguments
+      #   ensure_background_data_migration_succeeded("CopyColumn", arguments: ["users", "id", "id_for_type_change"])
+      #
+      def ensure_background_data_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.raise_in_prod_or_say_in_dev("Could not find background data migration for the given configuration: #{configuration}")
+        elsif !migration.succeeded?
+          raise "Expected background data migration for the given configuration to be marked as 'succeeded', " \
+                "but it is '#{migration.status}': #{configuration}"
+        end
+      end
+      alias ensure_background_migration_succeeded ensure_background_data_migration_succeeded
+
       # @private
       def create_background_data_migration(migration_name, *arguments, **options)
         options.assert_valid_keys(:batch_column_name, :min_value, :max_value, :batch_size, :sub_batch_size,

lib/online_migrations/background_schema_migrations/migration.rb

@@ -20,6 +20,7 @@ class Migration < ApplicationRecord
       self.table_name = :background_schema_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_schema_migrations/migration_helpers.rb

@@ -35,6 +35,29 @@ def remove_index_in_background(table_name, column_name = nil, name:, **options)
         enqueue_background_schema_migration(name, table_name, definition: definition, **migration_options)
       end
 
+      # Ensures that the background schema migration with the provided migration name succeeded.
+      #
+      # If the enqueued migration was not found in development (probably when resetting a dev environment
+      # followed by `db:migrate`), then a log warning is printed.
+      # If enqueued migration was not found in production, then the error is raised.
+      # If enqueued migration was found but is not succeeded, then the error is raised.
+      #
+      # @param migration_name [String, Symbol] Background schema migration name
+      #
+      # @example
+      #   ensure_background_schema_migration_succeeded("index_users_on_email")
+      #
+      def ensure_background_schema_migration_succeeded(migration_name)
+        migration = Migration.parents.find_by(migration_name: migration_name)
+
+        if migration.nil?
+          Utils.raise_in_prod_or_say_in_dev("Could not find background schema migration: '#{migration_name}'")
+        elsif !migration.succeeded?
+          raise "Expected background schema migration '#{migration_name}' to be marked as 'succeeded', " \
+                "but it is '#{migration.status}'."
+        end
+      end
+
       def enqueue_background_schema_migration(name, table_name, **options)
         if options[:connection_class_name].nil? && Utils.multiple_databases?
           raise ArgumentError, "You must pass a :connection_class_name when using multiple databases."

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

lib/online_migrations/utils.rb

@@ -40,6 +40,14 @@ def raise_or_say(message)
         end
       end
 
+      def raise_in_prod_or_say_in_dev(message)
+        if developer_env?
+          say(message)
+        else
+          raise message
+        end
+      end
+
       def warn(message)
         Kernel.warn("[online_migrations] #{message}")
       end

test/background_migrations/background_migrations.rb

@@ -2,6 +2,8 @@
 
 module BackgroundMigrations
   class MakeAllNonAdmins < OnlineMigrations::BackgroundMigration
+    def initialize(*_dummy_args) end
+
     def relation
       User.all
     end

test/schema_statements/misc_test.rb

@@ -36,6 +36,7 @@ def setup
 
     def teardown
       OnlineMigrations::BackgroundMigrations::Migration.delete_all
+      OnlineMigrations::BackgroundSchemaMigrations::Migration.delete_all
       @connection.drop_table(:users, if_exists: true)
       @connection.drop_table(:invoices, if_exists: true)
     end
@@ -224,6 +225,66 @@ def test_remove_background_data_migration_with_arguments
       assert_equal 0, OnlineMigrations::BackgroundMigrations::Migration.count
     end
 
+    def test_ensure_background_data_migration_succeeded
+      m = @connection.enqueue_background_data_migration("MakeAllNonAdmins")
+      assert m.succeeded?
+      assert_nothing_raised do
+        @connection.ensure_background_data_migration_succeeded("MakeAllNonAdmins")
+      end
+    end
+
+    def test_ensure_background_data_migration_succeeded_with_arguments
+      m = @connection.enqueue_background_data_migration("MakeAllNonAdmins", 1, { foo: "bar" }, 2)
+      assert m.succeeded?
+      assert_nothing_raised do
+        @connection.ensure_background_data_migration_succeeded("MakeAllNonAdmins", arguments: [1, { foo: "bar" }, 2])
+      end
+    end
+
+    def test_ensure_background_data_migration_succeeded_when_migration_not_found
+      out, = capture_io do
+        logger = ActiveSupport::Logger.new($stdout)
+        ActiveRecord::Base.stub(:logger, logger) do
+          @connection.ensure_background_data_migration_succeeded("MakeAllNonAdmins")
+        end
+      end
+      assert_match(/Could not find background data migration/i, out)
+    end
+
+    def test_ensure_background_data_migration_succeeded_when_migration_is_failed
+      m = @connection.enqueue_background_data_migration("MakeAllNonAdmins")
+      m.update_column(:status, :failed)
+
+      assert_raises_with_message(RuntimeError, /to be marked as 'succeeded'/i) do
+        @connection.ensure_background_data_migration_succeeded("MakeAllNonAdmins")
+      end
+    end
+
+    def test_ensure_background_schema_migration_succeeded
+      m = @connection.add_index_in_background(:users, :name, connection_class_name: "User")
+      assert m.succeeded?
+      assert_nothing_raised do
+        @connection.ensure_background_schema_migration_succeeded(m.name)
+      end
+    end
+
+    def test_ensure_background_schema_migration_succeeded_when_migration_not_found
+      Rails.stub(:env, ActiveSupport::StringInquirer.new("production")) do
+        assert_raises_with_message(RuntimeError, /Could not find background schema migration/i) do
+          @connection.ensure_background_schema_migration_succeeded("index_users_on_name")
+        end
+      end
+    end
+
+    def test_ensure_background_schema_migration_succeeded_when_migration_is_failed
+      m = @connection.add_index_in_background(:users, :name, connection_class_name: "User")
+      m.update_column(:status, :failed)
+
+      assert_raises_with_message(RuntimeError, /to be marked as 'succeeded'/i) do
+        @connection.ensure_background_schema_migration_succeeded("index_users_on_name")
+      end
+    end
+
     def test_disable_statement_timeout
       prev_value = get_statement_timeout
       set_statement_timeout(10)