Add ability to cancel background migrations

CHANGELOG.md

@@ -1,5 +1,7 @@
 ## master (unreleased)
 
+- Add ability to cancel background migrations
+
 ## 0.18.0 (2024-05-07)
 
 - Fix setting `started_at`/`finished_at` for parents of sharded background schema migrations

docs/background_data_migrations.md

@@ -237,6 +237,7 @@ Background Migrations can be in various states during its execution:
   Note: In normal circumstances, this should not be used since background migrations should be run and finished by the scheduler.
 * **failed**: A migration raises an exception when running.
 * **succeeded**: A migration finished without error.
+* **cancelled**: A migration was cancelled by the user.
 
 To get the progress (assuming `#count` method on background migration class was defined):
 
@@ -258,6 +259,15 @@ migration.retry # => `true` if scheduled to be retried, `false` - if not
 
 The migration will be retried on the next Scheduler run.
 
+## Cancelling a migration
+
+To cancel an existing migration from future performing, run:
+
+```ruby
+migration = OnlineMigrations::BackgroundMigrations::Migration.find(id)
+migration.cancel
+```
+
 ## Configuring
 
 There are a few configurable options for the Background Migrations. Custom configurations should be placed in a `online_migrations.rb` initializer.

docs/background_schema_migrations.md

@@ -77,6 +77,15 @@ migration.retry # => `true` if scheduled to be retried, `false` - if not
 
 The migration will be retried on the next Scheduler run.
 
+## Cancelling a migration
+
+To cancel an existing migration from future performing, run:
+
+```ruby
+migration = OnlineMigrations::BackgroundSchemaMigrations::Migration.find(id)
+migration.cancel
+```
+
 ## Instrumentation
 
 Background schema migrations use the [ActiveSupport::Notifications](http://api.rubyonrails.org/classes/ActiveSupport/Notifications.html) API.
@@ -121,6 +130,7 @@ Background Schema Migrations can be in various states during its execution:
 * **running**: A migration is being performed by a migration executor.
 * **failed**: A migration raises an exception when running.
 * **succeeded**: A migration finished without error.
+* **cancelled**: A migration was cancelled by the user.
 
 ## Configuring
 

lib/online_migrations/background_migrations/migration.rb

@@ -15,6 +15,7 @@ class Migration < ApplicationRecord
         :finishing,   # The migration is being manually finishing inline by the user.
         :failed,      # The migration raises an exception when running.
         :succeeded,   # The migration finished without error.
+        :cancelled,   # The migration was cancelled by the user.
       ]
 
       self.table_name = :background_migrations
@@ -98,6 +99,17 @@ def running!
         end
       end
 
+      # Overwrite enum's generated method to correctly work for composite migrations.
+      def cancelled!
+        return super if !composite?
+
+        transaction do
+          super
+          children.each { |child| child.cancelled! if !child.succeeded? }
+        end
+      end
+      alias cancel cancelled!
+
       def last_job
         migration_jobs.order(:max_value).last
       end

lib/online_migrations/background_migrations/migration_job.rb

@@ -8,6 +8,7 @@ class MigrationJob < ApplicationRecord
         :running,
         :failed,
         :succeeded,
+        :cancelled,
       ]
 
       self.table_name = :background_migration_jobs

lib/online_migrations/background_migrations/migration_job_status_validator.rb

@@ -5,9 +5,9 @@ module BackgroundMigrations
     # @private
     class MigrationJobStatusValidator < ActiveModel::Validator
       VALID_STATUS_TRANSITIONS = {
-        "enqueued" => ["running"],
-        "running" => ["succeeded", "failed"],
-        "failed" => ["enqueued", "running"],
+        "enqueued" => ["running", "cancelled"],
+        "running" => ["succeeded", "failed", "cancelled"],
+        "failed" => ["enqueued", "running", "cancelled"],
       }
 
       def validate(record)

lib/online_migrations/background_migrations/migration_runner.rb

@@ -13,6 +13,7 @@ def initialize(migration)
       # Runs one background migration job.
       def run_migration_job
         raise "Should not be called on a composite (with sharding) migration" if migration.composite?
+        return if migration.cancelled?
 
         mark_as_running if migration.enqueued?
         migration_payload = notifications_payload(migration)
@@ -56,7 +57,7 @@ def run_all_migration_jobs
           raise "This method is not intended for use in production environments"
         end
 
-        return if migration.completed?
+        return if migration.completed? || migration.cancelled?
 
         mark_as_running
 
@@ -77,7 +78,7 @@ def run_all_migration_jobs
       # Keep running until the migration is finished.
       #
       def finish
-        return if migration.completed?
+        return if migration.completed? || migration.cancelled?
 
         if migration.composite?
           migration.children.each do |child_migration|

lib/online_migrations/background_migrations/migration_status_validator.rb

@@ -7,7 +7,7 @@ class MigrationStatusValidator < ActiveModel::Validator
       VALID_STATUS_TRANSITIONS = {
         # enqueued -> running occurs when the migration starts performing.
         # enqueued -> paused occurs when the migration is paused before starting.
-        "enqueued" => ["running", "paused"],
+        "enqueued" => ["running", "paused", "cancelled"],
         # running -> paused occurs when a user pauses the migration as
         #   it's performing.
         # running -> finishing occurs when a user manually finishes the migration.
@@ -18,15 +18,16 @@ class MigrationStatusValidator < ActiveModel::Validator
           "finishing",
           "succeeded",
           "failed",
+          "cancelled",
         ],
         # finishing -> succeeded occurs when the migration completes successfully.
         # finishing -> failed occurs when the migration raises an exception when running.
-        "finishing" => ["succeeded", "failed"],
+        "finishing" => ["succeeded", "failed", "cancelled"],
         # paused -> running occurs when the migration is resumed after being paused.
-        "paused" => ["running"],
+        "paused" => ["running", "cancelled"],
         # failed -> enqueued occurs when the failed migration jobs are retried after being failed.
         # failed -> running occurs when the failed migration is retried.
-        "failed" => ["enqueued", "running"],
+        "failed" => ["enqueued", "running", "cancelled"],
       }
 
       def validate(record)

lib/online_migrations/background_schema_migrations/migration.rb

@@ -13,6 +13,7 @@ class Migration < ApplicationRecord
         :running,     # The migration is being performed by a migration executor.
         :failed,      # The migration raises an exception when running.
         :succeeded,   # The migration finished without error.
+        :cancelled,   # The migration was cancelled by the user.
       ]
 
       MAX_IDENTIFIER_LENGTH = 63
@@ -82,6 +83,17 @@ def completed?
         succeeded? || failed?
       end
 
+      # Overwrite enum's generated method to correctly work for composite migrations.
+      def cancelled!
+        return super if !composite?
+
+        transaction do
+          super
+          children.each { |child| child.cancelled! if !child.succeeded? }
+        end
+      end
+      alias cancel cancelled!
+
       # Returns the progress of the background schema migration.
       #
       # @return [Float] value in range from 0.0 to 100.0

lib/online_migrations/background_schema_migrations/migration_runner.rb

@@ -11,6 +11,8 @@ def initialize(migration)
       end
 
       def run
+        return if migration.cancelled?
+
         mark_as_running if migration.enqueued? || migration.failed?
 
         if migration.composite?

lib/online_migrations/background_schema_migrations/migration_status_validator.rb

@@ -6,13 +6,13 @@ module BackgroundSchemaMigrations
     class MigrationStatusValidator < ActiveModel::Validator
       VALID_STATUS_TRANSITIONS = {
         # enqueued -> running occurs when the migration starts performing.
-        "enqueued" => ["running"],
+        "enqueued" => ["running", "cancelled"],
         # running -> succeeded occurs when the migration completes successfully.
         # running -> failed occurs when the migration raises an exception when running and retry attempts exceeded.
-        "running" => ["succeeded", "failed"],
+        "running" => ["succeeded", "failed", "cancelled"],
         # failed -> enqueued occurs when the failed migration is enqueued to be retried.
         # failed -> running occurs when the failed migration is retried.
-        "failed" => ["enqueued", "running"],
+        "failed" => ["enqueued", "running", "cancelled"],
       }
 
       def validate(record)

test/background_migrations/migration_test.rb

@@ -131,6 +131,23 @@ def test_running_bang_pauses_composite_migration
       assert child2.running?
     end
 
+    def test_cancelled_bang_cancells_migration
+      m = create_migration
+      m.cancelled!
+      assert m.cancelled?
+    end
+
+    def test_cancelled_bang_cancells_composite_migration
+      m = create_migration(migration_name: "MakeAllDogsNice")
+      child1, child2 = m.children.to_a
+      run_all_migration_jobs(child1)
+      m.cancelled!
+
+      assert m.cancelled?
+      assert child1.reload.succeeded?
+      assert child2.reload.cancelled?
+    end
+
     def test_empty_relation
       m = create_migration(migration_name: "EmptyRelation")
       assert_equal 1, m.min_value