online_migrations 0.15.0 → 0.16.1

data/lib/online_migrations/background_schema_migrations/config.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundSchemaMigrations
+    # Class representing configuration options for background schema migrations.
+    class Config
+      # Maximum number of run attempts
+      #
+      # When attempts are exhausted, the migration is marked as failed.
+      # @return [Integer] defaults to 5
+      #
+      attr_accessor :max_attempts
+
+      # Statement timeout value used when running background schema migration.
+      #
+      # @return [Integer] defaults to 1 hour
+      #
+      attr_accessor :statement_timeout
+
+      # The callback to perform when an error occurs in the migration.
+      #
+      # @example
+      #   OnlineMigrations.config.background_schema_migrations.error_handler = ->(error, errored_migration) do
+      #     Bugsnag.notify(error) do |notification|
+      #       notification.add_metadata(:background_schema_migration, { name: errored_migration.name })
+      #     end
+      #   end
+      #
+      # @return [Proc] the callback to perform when an error occurs in the migration
+      #
+      attr_accessor :error_handler
+
+      def initialize
+        @max_attempts = 5
+        @statement_timeout = 1.hour
+        @error_handler = ->(error, errored_migration) {}
+      end
+    end
+  end
+end

data/lib/online_migrations/background_schema_migrations/migration.rb

@@ -0,0 +1,217 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundSchemaMigrations
+    # Class representing background schema migration.
+    #
+    # @note The records of this class should not be created manually, but via
+    #   `enqueue_background_schema_migration` helper inside migrations.
+    #
+    class Migration < ApplicationRecord
+      STATUSES = [
+        :enqueued,    # The migration has been enqueued by the user.
+        :running,     # The migration is being performed by a migration executor.
+        :failed,      # The migration raises an exception when running.
+        :succeeded,   # The migration finished without error.
+      ]
+
+      MAX_IDENTIFIER_LENGTH = 63
+
+      self.table_name = :background_schema_migrations
+
+      scope :queue_order, -> { order(created_at: :asc) }
+      scope :runnable, -> { where(composite: false) }
+      scope :active, -> { where(status: [statuses[:enqueued], statuses[:running]]) }
+      scope :except_succeeded, -> { where.not(status: :succeeded) }
+
+      scope :stuck, -> do
+        runnable.active.where(<<~SQL)
+          updated_at <= NOW() - interval '1 second' * (COALESCE(statement_timeout, 60*60*24) + 60*10)
+        SQL
+      end
+
+      scope :retriable, -> do
+        failed_retriable = runnable.failed.where("attempts < max_attempts")
+
+        stuck_sql             = connection.unprepared_statement { stuck.to_sql }
+        failed_retriable_sql  = connection.unprepared_statement { failed_retriable.to_sql }
+
+        from(Arel.sql(<<~SQL))
+          (
+            (#{failed_retriable_sql})
+            UNION
+            (#{stuck_sql})
+          ) AS #{table_name}
+        SQL
+      end
+
+      alias_attribute :name, :migration_name
+
+      # Avoid deprecation warnings.
+      if Utils.ar_version >= 7
+        enum :status, STATUSES.index_with(&:to_s)
+      else
+        enum status: STATUSES.index_with(&:to_s)
+      end
+
+      belongs_to :parent, class_name: name, optional: true
+      has_many :children, class_name: name, foreign_key: :parent_id
+
+      validates :table_name, presence: true, length: { maximum: MAX_IDENTIFIER_LENGTH }
+      validates :definition, presence: true
+      validates :migration_name, presence: true, uniqueness: {
+        scope: :shard,
+        message: ->(object, data) do
+          message = "(#{data[:value]}) has already been taken."
+          if object.index_addition?
+            message += " Consider enqueuing index creation with a different index name via a `:name` option."
+          end
+          message
+        end,
+      }
+
+      validate :validate_children_statuses, if: -> { composite? && status_changed? }
+      validate :validate_connection_class, if: :connection_class_name?
+      validate :validate_table_exists
+      validates_with MigrationStatusValidator, on: :update
+
+      before_validation :set_defaults
+
+      def completed?
+        succeeded? || failed?
+      end
+
+      # Returns the progress of the background schema migration.
+      #
+      # @return [Float] value in range from 0.0 to 100.0
+      #
+      def progress
+        if succeeded?
+          100.0
+        elsif composite?
+          progresses = children.map(&:progress)
+          (progresses.sum.to_f / progresses.size).round(2)
+        else
+          0.0
+        end
+      end
+
+      # Mark this migration as ready to be processed again.
+      #
+      # This is used to manually retrying failed migrations.
+      #
+      def retry
+        if composite?
+          children.failed.each(&:retry)
+        elsif failed?
+          update!(
+            status: self.class.statuses[:enqueued],
+            attempts: 0,
+            started_at: nil,
+            finished_at: nil,
+            error_class: nil,
+            error_message: nil,
+            backtrace: nil
+          )
+        end
+      end
+
+      def index_addition?
+        definition.match?(/create (unique )?index/i)
+      end
+
+      # @private
+      def connection_class
+        if connection_class_name && (klass = connection_class_name.safe_constantize)
+          Utils.find_connection_class(klass)
+        else
+          ActiveRecord::Base
+        end
+      end
+
+      # @private
+      def attempts_exceeded?
+        attempts >= max_attempts
+      end
+
+      # @private
+      def run
+        on_shard do
+          connection = connection_class.connection
+
+          connection.with_lock_retries do
+            statement_timeout = self.statement_timeout || OnlineMigrations.config.statement_timeout
+
+            with_statement_timeout(connection, statement_timeout) do
+              if index_addition?
+                index = connection.indexes(table_name).find { |i| i.name == name }
+                if index
+                  # Use index validity from https://github.com/rails/rails/pull/45160
+                  # when switching to ActiveRecord >= 7.1.
+                  schema = connection.send(:__schema_for_table, table_name)
+                  if connection.send(:__index_valid?, name, schema: schema)
+                    return
+                  else
+                    connection.remove_index(table_name, name: name)
+                  end
+                end
+              end
+
+              connection.execute(definition)
+            end
+          end
+        end
+      end
+
+      private
+        def validate_children_statuses
+          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
+          end
+        end
+
+        def validate_connection_class
+          klass = connection_class_name.safe_constantize
+          if !(klass < ActiveRecord::Base)
+            errors.add(:connection_class_name, "is not an ActiveRecord::Base child class")
+          end
+        end
+
+        def validate_table_exists
+          # Skip this validation if we have invalid connection class name.
+          return if errors.include?(:connection_class_name)
+
+          on_shard do
+            if !connection_class.connection.table_exists?(table_name)
+              errors.add(:table_name, "'#{table_name}' does not exist")
+            end
+          end
+        end
+
+        def set_defaults
+          config = ::OnlineMigrations.config.background_schema_migrations
+          self.max_attempts ||= config.max_attempts
+          self.statement_timeout ||= config.statement_timeout
+        end
+
+        def on_shard(&block)
+          shard = (self.shard || connection_class.default_shard).to_sym
+          connection_class.connected_to(shard: shard, role: :writing, &block)
+        end
+
+        def with_statement_timeout(connection, timeout)
+          return yield if timeout.nil?
+
+          prev_value = connection.select_value("SHOW statement_timeout")
+          connection.execute("SET statement_timeout TO #{connection.quote(timeout.in_milliseconds)}")
+          yield
+        ensure
+          connection.execute("SET statement_timeout TO #{connection.quote(prev_value)}")
+        end
+    end
+  end
+end

data/lib/online_migrations/background_schema_migrations/migration_helpers.rb

@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundSchemaMigrations
+    module MigrationHelpers
+      def add_index_in_background(table_name, column_name, **options)
+        migration_options = options.extract!(:max_attempts, :statement_timeout, :connection_class_name)
+
+        if index_exists?(table_name, column_name, **options)
+          Utils.say("Index creation was not enqueued because the index already exists.")
+          return
+        end
+
+        options[:algorithm] = :concurrently
+        index, algorithm, if_not_exists = add_index_options(table_name, column_name, **options)
+
+        create_index = ActiveRecord::ConnectionAdapters::CreateIndexDefinition.new(index, algorithm, if_not_exists)
+        schema_creation = ActiveRecord::ConnectionAdapters::PostgreSQL::SchemaCreation.new(self)
+        definition = schema_creation.accept(create_index)
+
+        enqueue_background_schema_migration(index.name, table_name, definition: definition, **migration_options)
+      end
+
+      def remove_index_in_background(table_name, column_name = nil, name:, **options)
+        raise ArgumentError, "Index name must be specified" if name.blank?
+
+        migration_options = options.extract!(:max_attempts, :statement_timeout, :connection_class_name)
+
+        if !index_exists?(table_name, column_name, **options, name: name)
+          Utils.say("Index deletion was not enqueued because the index does not exist.")
+          return
+        end
+
+        definition = "DROP INDEX CONCURRENTLY IF EXISTS #{quote_column_name(name)}"
+        enqueue_background_schema_migration(name, table_name, definition: definition, **migration_options)
+      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."
+        end
+
+        migration = create_background_schema_migration(name, table_name, **options)
+
+        run_inline = OnlineMigrations.config.run_background_migrations_inline
+        if run_inline && run_inline.call
+          runner = MigrationRunner.new(migration)
+          runner.run
+        end
+
+        migration
+      end
+
+      # @private
+      def create_background_schema_migration(migration_name, table_name, **options)
+        options.assert_valid_keys(:definition, :max_attempts, :statement_timeout, :connection_class_name)
+        migration = Migration.new(migration_name: migration_name, table_name: table_name, **options)
+
+        shards = Utils.shard_names(migration.connection_class)
+        if shards.size > 1
+          migration.children = shards.map do |shard|
+            child = migration.dup
+            child.shard = shard
+            child
+          end
+
+          migration.composite = true
+        end
+
+        # This will save all the records using a transaction.
+        migration.save!
+        migration
+      end
+    end
+  end
+end

data/lib/online_migrations/background_schema_migrations/migration_runner.rb

@@ -0,0 +1,110 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundSchemaMigrations
+    # Runs single background schema migration.
+    class MigrationRunner
+      attr_reader :migration
+
+      def initialize(migration)
+        @migration = migration
+      end
+
+      def run
+        mark_as_running if migration.enqueued? || migration.failed?
+
+        if migration.composite?
+          migration.children.each do |child_migration|
+            runner = self.class.new(child_migration)
+            runner.run
+          end
+        else
+          do_run
+        end
+      end
+
+      private
+        def mark_as_running
+          Migration.transaction do
+            migration.running!
+            migration.parent.running! if migration.parent
+          end
+        end
+
+        def do_run
+          migration_payload = notifications_payload(migration)
+
+          if migration.attempts == 0
+            ActiveSupport::Notifications.instrument("started.background_schema_migrations", migration_payload)
+          else
+            ActiveSupport::Notifications.instrument("retried.background_schema_migrations", migration_payload)
+          end
+
+          if should_throttle?
+            ActiveSupport::Notifications.instrument("throttled.background_schema_migrations", migration_payload)
+            return
+          end
+
+          migration.update!(
+            attempts: migration.attempts + 1,
+            status: :running,
+            started_at: Time.current,
+            finished_at: nil,
+            error_class: nil,
+            error_message: nil,
+            backtrace: nil
+          )
+
+          ActiveSupport::Notifications.instrument("run.background_schema_migrations", migration_payload) do
+            migration.run
+          end
+
+          migration.update!(status: :succeeded, finished_at: Time.current)
+
+          ActiveSupport::Notifications.instrument("completed.background_schema_migrations", migration_payload)
+
+          complete_parent_if_needed(migration) if migration.parent.present?
+        rescue Exception => e # rubocop:disable Lint/RescueException
+          backtrace_cleaner = ::OnlineMigrations.config.backtrace_cleaner
+
+          migration.update!(
+            status: :failed,
+            finished_at: Time.current,
+            error_class: e.class.name,
+            error_message: e.message,
+            backtrace: backtrace_cleaner ? backtrace_cleaner.clean(e.backtrace) : e.backtrace
+          )
+
+          ::OnlineMigrations.config.background_schema_migrations.error_handler.call(e, migration)
+        end
+
+        def should_throttle?
+          ::OnlineMigrations.config.throttler.call
+        end
+
+        def complete_parent_if_needed(migration)
+          parent = migration.parent
+          completed = false
+
+          parent.with_lock do
+            children = parent.children.select(:status)
+            if children.all?(&:succeeded?)
+              parent.succeeded!
+              completed = true
+            elsif children.any?(&:failed?)
+              parent.failed!
+              completed = true
+            end
+          end
+
+          if completed
+            ActiveSupport::Notifications.instrument("completed.background_migrations", notifications_payload(migration))
+          end
+        end
+
+        def notifications_payload(migration)
+          { background_schema_migration: migration }
+        end
+    end
+  end
+end

data/lib/online_migrations/background_schema_migrations/migration_status_validator.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundSchemaMigrations
+    # @private
+    class MigrationStatusValidator < ActiveModel::Validator
+      VALID_STATUS_TRANSITIONS = {
+        # enqueued -> running occurs when the migration starts performing.
+        "enqueued" => ["running"],
+        # 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"],
+        # failed -> enqueued occurs when the failed migration is enqueued to be retried.
+        # failed -> running occurs when the failed migration is retried.
+        "failed" => ["enqueued", "running"],
+      }
+
+      def validate(record)
+        return if !record.status_changed?
+
+        previous_status, new_status = record.status_change
+        valid_new_statuses = VALID_STATUS_TRANSITIONS.fetch(previous_status, [])
+
+        if !valid_new_statuses.include?(new_status)
+          record.errors.add(
+            :status,
+            "cannot transition background schema migration from status #{previous_status} to #{new_status}"
+          )
+        end
+      end
+    end
+  end
+end

data/lib/online_migrations/background_schema_migrations/scheduler.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundSchemaMigrations
+    # Class responsible for scheduling background schema migrations.
+    # It selects a single migration and runs it if there is no currently running migration.
+    #
+    # Scheduler should be configured to run periodically, for example, via cron.
+    # @example Run via whenever
+    #   # add this to schedule.rb
+    #   every 1.minute do
+    #     runner "OnlineMigrations.run_background_schema_migrations"
+    #   end
+    #
+    class Scheduler
+      def self.run
+        new.run
+      end
+
+      # Runs Scheduler
+      def run
+        migration = Migration.runnable.enqueued.queue_order.first || Migration.retriable.queue_order.first
+        if migration
+          runner = MigrationRunner.new(migration)
+          runner.run
+        end
+      end
+    end
+  end
+end

data/lib/online_migrations/config.rb

@@ -181,6 +181,26 @@ module OnlineMigrations
     #
     attr_accessor :run_background_migrations_inline
 
+    # Allows to throttle background data or schema migrations based on external signal (e.g. database health)
+    #
+    # It will be called before each run.
+    # If throttled, the current run will be retried next time.
+    #
+    # @return [Proc]
+    #
+    # @example
+    #   OnlineMigrations.config.throttler = -> { DatabaseStatus.unhealthy? }
+    #
+    attr_reader :throttler
+
+    # The Active Support backtrace cleaner that will be used to clean the
+    # backtrace of a background data or schema migration that errors.
+    #
+    # @return [ActiveSupport::BacktraceCleaner, nil] the backtrace cleaner to
+    #   use when cleaning a background migrations's backtrace. Defaults to `Rails.backtrace_cleaner`
+    #
+    attr_accessor :backtrace_cleaner
+
     # Configuration object to configure background migrations
     #
     # @return [BackgroundMigrationsConfig]
@@ -188,6 +208,8 @@ module OnlineMigrations
     #
     attr_reader :background_migrations
 
+    attr_reader :background_schema_migrations
+
     def initialize
       @table_renames = {}
       @column_renames = {}
@@ -202,6 +224,7 @@ module OnlineMigrations
       )
 
       @background_migrations = BackgroundMigrations::Config.new
+      @background_schema_migrations = BackgroundSchemaMigrations::Config.new
 
       @checks = []
       @start_after = 0
@@ -213,6 +236,7 @@ module OnlineMigrations
       @enabled_checks = @error_messages.keys.index_with({})
       @verbose_sql_logs = defined?(Rails.env) && (Rails.env.production? || Rails.env.staging?)
       @run_background_migrations_inline = -> { Utils.developer_env? }
+      @throttler = -> { false }
     end
 
     def lock_retrier=(value)
@@ -223,6 +247,14 @@ module OnlineMigrations
       @small_tables = table_names.map(&:to_s)
     end
 
+    def throttler=(value)
+      if !value.respond_to?(:call)
+        raise ArgumentError, "throttler must be a callable."
+      end
+
+      @throttler = value
+    end
+
     # Enables specific check
     #
     # For the list of available checks look at the `error_messages.rb` file.

data/lib/online_migrations/schema_statements.rb

@@ -4,6 +4,7 @@ module OnlineMigrations
   module SchemaStatements
     include ChangeColumnTypeHelpers
     include BackgroundMigrations::MigrationHelpers
+    include BackgroundSchemaMigrations::MigrationHelpers
 
     # Updates the value of a column in batches.
     #

data/lib/online_migrations/utils.rb

@@ -135,6 +135,13 @@ module OnlineMigrations
         connection.select_value(query) == "v"
       end
 
+      def find_connection_class(model)
+        model.ancestors.find do |parent|
+          parent == ActiveRecord::Base ||
+            (parent.is_a?(Class) && parent.abstract_class?)
+        end
+      end
+
       def shard_names(model)
         model.ancestors.each do |ancestor|
           # There is no official method to get shard names from the model.

data/lib/online_migrations/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module OnlineMigrations
-  VERSION = "0.15.0"
+  VERSION = "0.16.1"
 end

data/lib/online_migrations.rb

@@ -4,6 +4,7 @@ require "active_record"
 
 require "online_migrations/version"
 require "online_migrations/utils"
+require "online_migrations/background_schema_migrations/migration_helpers"
 require "online_migrations/change_column_type_helpers"
 require "online_migrations/background_migrations/migration_helpers"
 require "online_migrations/schema_statements"
@@ -22,6 +23,7 @@ module OnlineMigrations
 
   extend ActiveSupport::Autoload
 
+  autoload :ApplicationRecord
   autoload :BatchIterator
   autoload :VerboseSqlLogs
   autoload :ForeignKeysCollector
@@ -52,7 +54,6 @@ module OnlineMigrations
     autoload :DeleteOrphanedRecords
     autoload :PerformActionOnRelation
     autoload :ResetCounters
-    autoload :ApplicationRecord
     autoload :MigrationJob
     autoload :Migration
     autoload :MigrationJobRunner
@@ -60,6 +61,16 @@ module OnlineMigrations
     autoload :Scheduler
   end
 
+  module BackgroundSchemaMigrations
+    extend ActiveSupport::Autoload
+
+    autoload :Config
+    autoload :Migration
+    autoload :MigrationStatusValidator
+    autoload :MigrationRunner
+    autoload :Scheduler
+  end
+
   class << self
     # @private
     attr_accessor :current_migration
@@ -72,10 +83,16 @@ module OnlineMigrations
       @config ||= Config.new
     end
 
-    # Run background migrations
+    # Run background data migrations
     def run_background_migrations
       BackgroundMigrations::Scheduler.run
     end
+    alias run_background_data_migrations run_background_migrations
+
+    # Run background schema migrations
+    def run_background_schema_migrations
+      BackgroundSchemaMigrations::Scheduler.run
+    end
 
     def deprecator
       @deprecator ||=