Support sharding for background migrations

.github/workflows/test.yml

@@ -53,5 +53,13 @@ jobs:
         with:
           ruby-version: ${{ matrix.ruby-version }}
           bundler-cache: true
+      - name: Prepare PostgreSQL shards
+        run: |
+          createdb online_migrations_shard_one
+          createdb online_migrations_shard_two
+        env:
+          PGHOST: localhost
+          PGUSER: postgres
+          PGPASSWORD: postgres
       - name: Run the test suite
         run: bundle exec rake test

.rubocop.yml

@@ -124,6 +124,7 @@ Naming/FileName:
   Exclude:
     - lib/online_migrations/version.rb
     - test/support/schema.rb
+    - test/support/db/**
     - test/test_helper.rb
     - gemfiles/**.gemfile
 

CHANGELOG.md

@@ -1,5 +1,17 @@
 ## master (unreleased)
 
+- Support sharding for background migrations
+
+    Now, if a `relation` inside background migration definition is defined on a sharded model,
+    then that background migration would automatically run on all the shards.
+
+    To get all the new sharding related schema changes, you need to run:
+
+    ```sh
+    $ bin/rails generate online_migrations:upgrade
+    $ bin/rails db:migrate
+    ```
+
 - Drop support for Ruby < 2.7 and Rails < 6.1
 
 ## 0.10.0 (2023-12-12)

README.md

@@ -37,10 +37,20 @@ And then run:
 ```sh
 $ bundle install
 $ bin/rails generate online_migrations:install
+$ bin/rails db:migrate
 ```
 
 **Note**: If you do not have plans on using [background migrations](docs/background_migrations.md) feature, then you can delete the generated migration and regenerate it later, if needed.
 
+### Upgrading
+
+If you're already using [background migrations](docs/background_migrations.md), your background migrations tables may require additional columns. After every upgrade, please run:
+
+```sh
+$ bin/rails generate online_migrations:upgrade
+$ bin/rails db:migrate
+```
+
 ## Motivation
 
 Writing a safe migration can be daunting. Numerous articles have been written on the topic and a few gems are trying to address the problem. Even for someone who has a pretty good command of PostgreSQL, remembering all the subtleties of explicit locking can be problematic.

docs/background_migrations.md

@@ -116,13 +116,15 @@ enqueue_background_migration("MyMigrationWithArgs", arg1, arg2, ...)
 
 ## Predefined background migrations
 
-* `BackfillColumn` - backfills column(s) with scalar values (enqueue using `backfill_column_in_background`)
+* `BackfillColumn` - backfills column(s) with scalar values (enqueue using `backfill_column_in_background`; or `backfill_column_for_type_change_in_background` if backfilling column for which type change is in progress)
 * `CopyColumn` - copies data from one column(s) to other(s) (enqueue using `copy_column_in_background`)
 * `DeleteAssociatedRecords` - deletes records associated with a parent object (enqueue using `delete_associated_records_in_background`)
 * `DeleteOrphanedRecords` - deletes records with one or more missing relations (enqueue using `delete_orphaned_records_in_background`)
 * `PerformActionOnRelation` - performs specific action on a relation or individual records (enqueue using `perform_action_on_relation_in_background`)
 * `ResetCounters` - resets one or more counter caches to their correct value (enqueue using `reset_counters_in_background`)
 
+**Note**: These migration helpers should be run inside the migration against the database where background migrations tables are defined.
+
 ## 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/generators/online_migrations/templates/add_sharding_to_online_migrations.rb.tt

@@ -0,0 +1,18 @@
+class AddShardingToOnlineMigrations < <%= migration_parent %>
+  def change
+    safety_assured do
+      remove_index :background_migrations, [:migration_name, :arguments], unique: true
+
+      change_table :background_migrations do |t|
+        t.bigint :parent_id
+        t.string :shard
+        t.boolean :runnable, default: true, null: false
+
+        t.foreign_key :background_migrations, column: :parent_id, on_delete: :cascade
+
+        t.index [:migration_name, :arguments, :shard],
+          unique: true, name: :index_background_migrations_on_unique_configuration
+      end
+    end
+  end
+end

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

@@ -1,6 +1,7 @@
 class InstallOnlineMigrations < <%= migration_parent %>
   def change
     create_table :background_migrations do |t|
+      t.bigint :parent_id
       t.string :migration_name, null: false
       t.jsonb :arguments, default: [], null: false
       t.string :batch_column_name, null: false
@@ -13,9 +14,13 @@ class InstallOnlineMigrations < <%= migration_parent %>
       t.integer :sub_batch_pause_ms, null: false
       t.integer :batch_max_attempts, null: false
       t.string :status, default: "enqueued", null: false
+      t.string :shard
+      t.boolean :runnable, default: true, null: false
       t.timestamps null: false
 
-      t.index [:migration_name, :arguments],
+      t.foreign_key :background_migrations, column: :parent_id, on_delete: :cascade
+
+      t.index [:migration_name, :arguments, :shard],
         unique: true, name: :index_background_migrations_on_unique_configuration
     end
 

lib/generators/online_migrations/upgrade_generator.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record/migration"
+
+module OnlineMigrations
+  # @private
+  class UpgradeGenerator < Rails::Generators::Base
+    include ActiveRecord::Generators::Migration
+
+    source_root File.expand_path("templates", __dir__)
+
+    def copy_templates
+      migrations_to_be_applied.each do |migration|
+        migration_template("#{migration}.rb", File.join(db_migrate_path, "#{migration}.rb"))
+      end
+    end
+
+    private
+      def migrations_to_be_applied
+        connection = BackgroundMigrations::Migration.connection
+        columns = connection.columns(BackgroundMigrations::Migration.table_name).map(&:name)
+
+        migrations = []
+        migrations << "add_sharding_to_online_migrations" if !columns.include?("shard")
+        migrations
+      end
+
+      def migration_parent
+        "ActiveRecord::Migration[#{Utils.ar_version}]"
+      end
+  end
+end

lib/online_migrations/background_migrations/migration.rb

@@ -15,14 +15,18 @@ class Migration < ActiveRecord::Base
       self.table_name = :background_migrations
 
       scope :queue_order, -> { order(created_at: :asc) }
+      scope :runnable, -> { where(runnable: true) }
       scope :active, -> { where(status: [statuses[:enqueued], statuses[:running]]) }
+      scope :except_succeeded, -> { where.not(status: :succeeded) }
       scope :for_migration_name, ->(migration_name) { where(migration_name: normalize_migration_name(migration_name)) }
       scope :for_configuration, ->(migration_name, arguments) do
         for_migration_name(migration_name).where("arguments = ?", arguments.to_json)
       end
 
       enum status: STATUSES.index_with(&:to_s)
 
+      belongs_to :parent, class_name: name, optional: true
+      has_many :children, class_name: name, foreign_key: :parent_id
       has_many :migration_jobs
 
       validates :migration_name, :batch_column_name, presence: true
@@ -33,7 +37,7 @@ class Migration < ActiveRecord::Base
       validates :batch_pause, :sub_batch_pause_ms, presence: true,
                   numericality: { greater_than_or_equal_to: 0 }
       validates :rows_count, numericality: { greater_than_or_equal_to: 0 }, allow_nil: true
-      validates :arguments, uniqueness: { scope: :migration_name }
+      validates :arguments, uniqueness: { scope: [:migration_name, :shard] }
 
       validate :validate_batch_column_values
       validate :validate_batch_sizes
@@ -43,6 +47,7 @@ class Migration < ActiveRecord::Base
       validates_with MigrationStatusValidator, on: :update
 
       before_validation :set_defaults
+      before_create :create_child_migrations, unless: :parent_id?
 
       # @private
       def self.normalize_migration_name(migration_name)
@@ -58,12 +63,36 @@ def completed?
         succeeded? || failed?
       end
 
+      def composite?
+        parent_id.nil? && !runnable?
+      end
+
+      # Overwrite enum's generated method to correctly work for composite migrations.
+      def paused!
+        return super if !composite?
+
+        transaction do
+          super
+          children.each { |child| child.paused! if child.enqueued? || child.running? }
+        end
+      end
+
+      # Overwrite enum's generated method to correctly work for composite migrations.
+      def running!
+        return super if !composite?
+
+        transaction do
+          super
+          children.each { |child| child.running! if child.paused? }
+        end
+      end
+
       def last_job
-        migration_jobs.order(max_value: :desc).first
+        migration_jobs.order(:max_value).last
       end
 
       def last_completed_job
-        migration_jobs.completed.order(finished_at: :desc).first
+        migration_jobs.completed.order(:finished_at).last
       end
 
       # Returns the progress of the background migration.
@@ -75,6 +104,11 @@ def last_completed_job
       def progress
         if succeeded?
           1.0
+        elsif composite?
+          progresses = children.map(&:progress).compact
+          if progresses.any?
+            progresses.sum / progresses.size
+          end
         elsif rows_count
           jobs_rows_count = migration_jobs.succeeded.sum(:batch_size)
           # The last migration job may need to process the amount of rows
@@ -95,6 +129,10 @@ def migration_relation
         migration_object.relation
       end
 
+      def migration_model
+        migration_relation.model
+      end
+
       # Returns whether the interval between previous step run has passed.
       # @return [Boolean]
       #
@@ -170,7 +208,13 @@ def validate_batch_sizes
         end
 
         def validate_jobs_status
-          if succeeded? && migration_jobs.except_succeeded.exists?
+          if composite?
+            if succeeded? && children.except_succeeded.exists?
+              errors.add(:base, "all child migrations must be succeeded")
+            elsif failed? && !children.failed.exists?
+              errors.add(:base, "at least one child migration must be failed")
+            end
+          elsif succeeded? && migration_jobs.except_succeeded.exists?
             errors.add(:base, "all migration jobs must be succeeded")
           elsif failed? && !migration_jobs.failed.exists?
             errors.add(:base, "at least one migration job must be failed")
@@ -201,6 +245,20 @@ def set_defaults
           end
         end
 
+        def create_child_migrations
+          shards = Utils.shard_names(migration_model)
+
+          if shards.size > 1
+            children = shards.map do |shard|
+              child = dup
+              child.shard = shard
+              child
+            end
+            self.runnable = false
+            self.children = children
+          end
+        end
+
         def next_min_value
           if last_job
             last_job.max_value.next