Add remove_background_migration migration helper

CHANGELOG.md

@@ -1,5 +1,6 @@
 ## master (unreleased)
 
+- Add `remove_background_migration` migration helper
 - Allow adding bigint foreign keys referencing integer primary keys
 - Fix `add_reference_concurrently` to check for mismatched key types
 

docs/background_migrations.md

@@ -30,7 +30,8 @@ A generator is provided to create background migrations. Generate a new backgrou
 $ bin/rails generate online_migrations:background_migration backfill_project_issues_count
 ```
 
-This creates the background migration file `lib/online_migrations/background_migrations/backfill_project_issues_count.rb`.
+This creates the background migration file `lib/online_migrations/background_migrations/backfill_project_issues_count.rb`
+and the regular migration file `db/migrate/xxxxxxxxxxxxxx_enqueue_backfill_project_issues_count.rb` where we enqueue it.
 
 The generated class is a subclass of `OnlineMigrations::BackgroundMigration` that implements:
 
@@ -76,12 +77,16 @@ end
 You can enqueue your background migration to be run by the scheduler via:
 
 ```ruby
-# db/migrate/xxxxxxxxxxxxxx_backfill_project_issues_count.rb
-# ...
-def up
-  enqueue_background_migration("BackfillProjectIssuesCount")
+# db/migrate/xxxxxxxxxxxxxx_enqueue_backfill_project_issues_count.rb
+class EnqueueBackfillProjectIssuesCount < ActiveRecord::Migration[7.1]
+  def up
+    enqueue_background_migration("BackfillProjectIssuesCount")
+  end
+
+  def down
+    remove_background_migration("BackfillProjectIssuesCount")
+  end
 end
-# ...
 ```
 
 `enqueue_background_migration` accepts additional configuration options which controls how the background migration is run. Check the [source code](https://github.com/fatkodima/online_migrations/blob/master/lib/online_migrations/background_migrations/migration_helpers.rb) for the list of all available configuration options.
@@ -106,7 +111,17 @@ end
 And pass them when enqueuing:
 
 ```ruby
-enqueue_background_migration("MyMigrationWithArgs", arg1, arg2, ...)
+def up
+  enqueue_background_migration("MyMigrationWithArgs", arg1, arg2, ...)
+end
+```
+
+Make sure to also pass the arguments inside the `down` method of the migration:
+
+```ruby
+def down
+  remove_background_migration("MyMigrationWithArgs", arg1, arg2, ...)
+end
 ```
 
 ## Considerations when writing Background Migrations

lib/generators/online_migrations/background_migration_generator.rb

@@ -1,12 +1,15 @@
 # frozen_string_literal: true
 
 require "rails/generators"
+require "rails/generators/active_record/migration"
 
 module OnlineMigrations
   # @private
   class BackgroundMigrationGenerator < Rails::Generators::NamedBase
+    include ActiveRecord::Generators::Migration
+
     source_root File.expand_path("templates", __dir__)
-    desc "This generator creates a background migration file."
+    desc "This generator creates a background migration related files."
 
     def create_background_migration_file
       migrations_module_file_path = migrations_module.underscore
@@ -20,6 +23,10 @@ def create_background_migration_file
       template("background_migration.rb", template_file)
     end
 
+    def create_migration_file
+      migration_template("migration.rb", File.join(db_migrate_path, "enqueue_#{file_name}.rb"))
+    end
+
     private
       def migrations_module
         config.migrations_module
@@ -28,5 +35,9 @@ def migrations_module
       def config
         OnlineMigrations.config.background_migrations
       end
+
+      def migration_parent
+        "ActiveRecord::Migration[#{Utils.ar_version}]"
+      end
   end
 end

lib/generators/online_migrations/templates/migration.rb.tt

@@ -0,0 +1,10 @@
+class Enqueue<%= class_name %> < <%= migration_parent %>
+  def up
+    enqueue_background_migration("<%= class_name %>", ...args)
+  end
+
+  def down
+    # Make sure to pass the same arguments as in the "up" method, if any.
+    remove_background_migration("<%= class_name %>", ...args)
+  end
+end

lib/online_migrations/background_migrations/migration_helpers.rb

@@ -381,11 +381,26 @@ def enqueue_background_migration(migration_name, *arguments, **options)
         migration
       end
 
+      # Removes the background migration for the given class name and arguments, if exists.
+      #
+      # @param migration_name [String, Class] Background migration job class name
+      # @param arguments [Array] Extra arguments the migration was originally created with
+      #
+      # @example
+      #   remove_background_migration("BackfillProjectIssuesCount")
+      #
+      def remove_background_migration(migration_name, *arguments)
+        migration_name = migration_name.name if migration_name.is_a?(Class)
+        Migration.for_configuration(migration_name, arguments).delete_all
+      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.new(
           migration_name: migration_name,
           arguments: arguments,

lib/online_migrations/command_recorder.rb

@@ -26,6 +26,7 @@ module CommandRecorder
       :add_reference_concurrently,
       :change_column_type_in_background,
       :enqueue_background_migration,
+      :remove_background_migration,
 
       # column type change helpers
       :initialize_column_type_change,

test/background_migration_generator_test.rb

@@ -21,6 +21,16 @@ def test_creates_background_migration_file
     end
   end
 
+  def test_creates_migration_file
+    run_generator(["make_all_non_admins"])
+
+    assert_migration("db/migrate/enqueue_make_all_non_admins.rb") do |content|
+      assert_includes content, "class EnqueueMakeAllNonAdmins < ActiveRecord::Migration"
+      assert_includes content, 'enqueue_background_migration("MakeAllNonAdmins"'
+      assert_includes content, 'remove_background_migration("MakeAllNonAdmins"'
+    end
+  end
+
   def test_generator_uses_configured_migrations_path
     OnlineMigrations.config.background_migrations.stub(:migrations_path, "app/lib") do
       run_generator(["make_all_non_admins"])

test/background_migrations/background_migrations.rb

@@ -121,6 +121,20 @@ def count
     end
   end
 
+  class MigrationWithArguments < OnlineMigrations::BackgroundMigration
+    def initialize(_arg1, _arg2)
+      # no-op
+    end
+
+    def relation
+      User.none
+    end
+
+    def process_batch(_users)
+      # no-op
+    end
+  end
+
   class NotAMigration
     def relation; end
     def process_batch(*); end

test/schema_statements/misc_test.rb

@@ -203,6 +203,24 @@ def test_run_background_migrations_inline_configured_to_custom_proc
       assert_nil user.reload.admin
     end
 
+    def test_remove_non_existing_background_migration
+      assert_nothing_raised do
+        @connection.remove_background_migration("NonExistent")
+      end
+    end
+
+    def test_remove_background_migration
+      @connection.enqueue_background_migration("MakeAllNonAdmins")
+      @connection.remove_background_migration("MakeAllNonAdmins")
+      assert_equal 0, OnlineMigrations::BackgroundMigrations::Migration.count
+    end
+
+    def test_remove_background_migration_with_arguments
+      @connection.enqueue_background_migration("MigrationWithArguments", 1, { a: 2 })
+      @connection.remove_background_migration("MigrationWithArguments", 1, { a: 2 })
+      assert_equal 0, OnlineMigrations::BackgroundMigrations::Migration.count
+    end
+
     def test_disable_statement_timeout
       prev_value = get_statement_timeout
       set_statement_timeout(10)