online_migrations 0.26.0 → 0.27.0

data/lib/online_migrations/background_schema_migrations/migration_helpers.rb

@@ -78,67 +78,46 @@ module OnlineMigrations
       #   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)
+        migrations = Migration.where(migration_name: migration_name).to_a
 
-        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}'."
+        if migrations.empty?
+          Utils.raise_in_prod_or_say_in_dev("Could not find background schema migration(s): '#{migration_name}'.")
+        elsif !migrations.all?(&:succeeded?)
+          raise "Expected background schema migration(s) '#{migration_name}' to be marked as 'succeeded'."
         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."
-        end
-
-        migration = create_background_schema_migration(name, table_name, **options)
+      def enqueue_background_schema_migration(migration_name, table_name, connection_class_name: nil, **options)
+        options.assert_valid_keys(:definition, :max_attempts, :statement_timeout)
 
-        run_inline = OnlineMigrations.config.run_background_migrations_inline
-        if run_inline && run_inline.call
-          runner = MigrationRunner.new(migration)
-          runner.run
+        if Utils.multiple_databases? && !connection_class_name
+          raise ArgumentError, "You must pass a :connection_class_name when using multiple databases."
         end
 
-        migration
-      end
-
-      # @private
-      def create_background_schema_migration(migration_name, table_name, connection_class_name: nil, **options)
-        options.assert_valid_keys(:definition, :max_attempts, :statement_timeout)
-
         if connection_class_name
-          connection_class_name = __normalize_connection_class_name(connection_class_name)
+          klass = connection_class_name.constantize
+          connection_class = Utils.find_connection_class(klass)
+          # Normalize to the real connection class name.
+          connection_class_name = connection_class.name
+        else
+          connection_class = ActiveRecord::Base
         end
 
-        Migration.find_or_create_by!(migration_name: migration_name, shard: nil,
-                                     connection_class_name: connection_class_name) do |migration|
-          migration.assign_attributes(**options, table_name: table_name)
+        shards = Utils.shard_names(connection_class)
+        shards = [nil] if shards.size == 1
 
-          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
+        shards.each do |shard|
+          migration = Migration.create_with(**options, table_name: table_name)
+                               .find_or_create_by!(migration_name: migration_name, shard: shard, connection_class_name: connection_class_name)
 
-            migration.composite = true
+          if Utils.run_background_migrations_inline?
+            runner = MigrationRunner.new(migration)
+            runner.run
           end
         end
-      end
 
-      private
-        def __normalize_connection_class_name(connection_class_name)
-          if connection_class_name
-            klass = connection_class_name.safe_constantize
-            if klass
-              connection_class = Utils.find_connection_class(klass)
-              connection_class.name if connection_class
-            end
-          end
-        end
+        true
+      end
     end
   end
 end

data/lib/online_migrations/background_schema_migrations/migration_runner.rb

@@ -13,119 +13,65 @@ module OnlineMigrations
       def run
         return if migration.cancelled? || migration.succeeded?
 
-        mark_as_running if migration.enqueued? || migration.errored?
+        migration.running! if migration.enqueued? || migration.errored?
+        migration_payload = { migration: migration }
 
-        if migration.composite?
-          migration.children.each do |child_migration|
-            runner = self.class.new(child_migration)
-            runner.run
-          end
+        if migration.attempts == 0
+          ActiveSupport::Notifications.instrument("started.background_schema_migrations", migration_payload)
         else
-          do_run
+          ActiveSupport::Notifications.instrument("retried.background_schema_migrations", migration_payload)
         end
-      end
-
-      private
-        def mark_as_running
-          Migration.transaction do
-            migration.running!
 
-            if (parent = migration.parent)
-              if parent.started_at
-                parent.update!(status: :running, finished_at: nil)
-              else
-                parent.update!(status: :running, started_at: Time.current, finished_at: nil)
-              end
-            end
-          end
+        if should_throttle?
+          ActiveSupport::Notifications.instrument("throttled.background_schema_migrations", migration_payload)
+          return
         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
-
-          # Background schema migrations could take a while to run. It is possible, that the process
-          # never reaches this (or the rescue below) line of code. E.g., when it is force quitted
-          # (SIGKILL etc.) and so the migration will end up in the "running" state and the query is
-          # still executing (or already finished) in the database. This migration can either be safely
-          # manually retried or will be picked up in the future by scheduler when it decides that
-          # this migration is stuck.
-
-          migration.update!(status: :succeeded, finished_at: Time.current)
+        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
 
-          ActiveSupport::Notifications.instrument("completed.background_schema_migrations", migration_payload)
+        # Background schema migrations could take a while to run. It is possible, that the process
+        # never reaches this (or the rescue below) line of code. E.g., when it is force quitted
+        # (SIGKILL etc.) and so the migration will end up in the "running" state and the query is
+        # still executing (or already finished) in the database. This migration can either be safely
+        # manually retried or will be picked up in the future by scheduler when it decides that
+        # this migration is stuck.
 
-          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: :succeeded, finished_at: Time.current)
 
-          status = migration.attempts_exceeded? ? :failed : :errored
+        ActiveSupport::Notifications.instrument("completed.background_schema_migrations", migration_payload)
+      rescue Exception => e # rubocop:disable Lint/RescueException
+        backtrace_cleaner = ::OnlineMigrations.config.backtrace_cleaner
 
-          migration.update!(
-            status: status,
-            finished_at: Time.current,
-            error_class: e.class.name,
-            error_message: e.message,
-            backtrace: backtrace_cleaner ? backtrace_cleaner.clean(e.backtrace) : e.backtrace
-          )
+        status = migration.attempts_exceeded? ? :failed : :errored
 
-          complete_parent_if_needed(migration) if migration.parent.present?
+        migration.update!(
+          status: status,
+          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)
-          raise if Utils.run_background_migrations_inline?
-        end
+        ::OnlineMigrations.config.background_schema_migrations.error_handler.call(e, migration)
+        raise if Utils.run_background_migrations_inline?
+      end
 
+      private
         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.update!(status: :succeeded, finished_at: Time.current)
-              completed = true
-            elsif children.any?(&:failed?)
-              parent.update!(status: :failed, finished_at: Time.current)
-              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/scheduler.rb

@@ -35,8 +35,8 @@ module OnlineMigrations
 
       private
         def find_migration(**options)
-          active_migrations = Migration.running.reject(&:stuck?)
-          runnable_migrations = Migration.runnable.enqueued.queue_order.to_a + Migration.retriable.queue_order.to_a
+          stuck_migrations, active_migrations = Migration.running.partition(&:stuck?)
+          runnable_migrations = (Migration.enqueued + Migration.errored + stuck_migrations).sort_by(&:created_at)
 
           if options.key?(:shard)
             runnable_migrations = runnable_migrations.select { |migration| migration.shard.to_s == options[:shard].to_s }

data/lib/online_migrations/change_column_type_helpers.rb

@@ -477,10 +477,11 @@ module OnlineMigrations
       end
 
       def __copy_check_constraints(table_name, from_column, to_column)
-        check_constraints = check_constraints(table_name).select { |c| c.expression.include?(from_column) }
+        from_column_re = /\b#{from_column}\b/
+        check_constraints = check_constraints(table_name).select { |c| c.expression.match?(from_column_re) }
 
         check_constraints.each do |check|
-          new_expression = check.expression.gsub(from_column, to_column)
+          new_expression = check.expression.gsub(from_column_re, to_column)
 
           add_check_constraint(table_name, new_expression, validate: false)
 

data/lib/online_migrations/config.rb

@@ -210,10 +210,10 @@ module OnlineMigrations
 
     # Configuration object to configure background migrations
     #
-    # @return [BackgroundMigrationsConfig]
-    # @see BackgroundMigrationsConfig
+    # @return [BackgroundDataMigrations::Config]
+    # @see BackgroundDataMigrations::Config
     #
-    attr_reader :background_migrations
+    attr_reader :background_data_migrations
 
     attr_reader :background_schema_migrations
 
@@ -230,7 +230,7 @@ module OnlineMigrations
         lock_timeout: 0.2.seconds
       )
 
-      @background_migrations = BackgroundMigrations::Config.new
+      @background_data_migrations = BackgroundDataMigrations::Config.new
       @background_schema_migrations = BackgroundSchemaMigrations::Config.new
 
       @checks = []

data/lib/online_migrations/data_migration.rb

@@ -0,0 +1,127 @@
+# frozen_string_literal: true
+
+gem "sidekiq", ">= 7.3.3"
+require "sidekiq"
+
+module OnlineMigrations
+  # Base class that is inherited by the host application's data migration classes.
+  class DataMigration
+    class NotFoundError < NameError; end
+
+    class << self
+      # Finds a Data Migration with the given name.
+      #
+      # @param name [String] the name of the Data Migration to be found.
+      #
+      # @return [DataMigration] the Data Migration with the given name.
+      #
+      # @raise [NotFoundError] if a Data Migration with the given name does not exist.
+      #
+      def named(name)
+        namespace = OnlineMigrations.config.background_data_migrations.migrations_module.constantize
+        internal_namespace = ::OnlineMigrations::BackgroundDataMigrations
+
+        migration = "#{namespace}::#{name}".safe_constantize ||
+                    "#{internal_namespace}::#{name}".safe_constantize
+
+        raise NotFoundError.new("Data Migration #{name} not found", name) if migration.nil?
+        if !(migration.is_a?(Class) && migration < self)
+          raise NotFoundError.new("#{name} is not a Data Migration", name)
+        end
+
+        migration
+      end
+
+      # @private
+      attr_accessor :active_record_enumerator_batch_size
+
+      # Limit the number of records that will be fetched in a single query when
+      # iterating over an Active Record collection migration.
+      #
+      # @param size [Integer] the number of records to fetch in a single query.
+      #
+      def collection_batch_size(size)
+        self.active_record_enumerator_batch_size = size
+      end
+    end
+
+    # A hook to override that will be called when the migration starts running.
+    #
+    def after_start
+    end
+
+    # A hook to override that will be called around 'process' each time.
+    #
+    # Can be useful for some metrics collection, performance tracking etc.
+    #
+    def around_process
+      yield
+    end
+
+    # A hook to override that will be called when the migration resumes its work.
+    #
+    def after_resume
+    end
+
+    # A hook to override that will be called each time the migration is interrupted.
+    #
+    # This can be due to interruption or sidekiq stopping.
+    #
+    def after_stop
+    end
+
+    # A hook to override that will be called when the migration finished its work.
+    #
+    def after_complete
+    end
+
+    # A hook to override that will be called when the migration is paused.
+    #
+    def after_pause
+    end
+
+    # A hook to override that will be called when the migration is cancelled.
+    #
+    def after_cancel
+    end
+
+    # The collection to be processed.
+    #
+    # @return [ActiveRecord::Relation, ActiveRecord::Batches::BatchEnumerator, Array, Enumerator]
+    #
+    # @raise [NotImplementedError] with a message advising subclasses to override this method.
+    #
+    def collection
+      raise NotImplementedError, "#{self.class.name} must implement a 'collection' method"
+    end
+
+    # The action to be performed on each item from the collection.
+    #
+    # @param _item the current item from the collection being iterated
+    # @raise [NotImplementedError] with a message advising subclasses to override this method.
+    #
+    def process(_item)
+      raise NotImplementedError, "#{self.class.name} must implement a 'process' method"
+    end
+
+    # Total count of iterations to be performed (optional, to be able to show progress).
+    #
+    # @return [Integer, nil]
+    #
+    def count
+    end
+
+    # Enumerator builder. You may override this method to return any Enumerator yielding
+    # pairs of `[item, item_cursor]`, instead of using `collection`.
+    #
+    # It is useful when it is not practical or impossible to define an explicit collection
+    # in the `collection` method.
+    #
+    # @param cursor [Object, nil] cursor position to resume from, or nil on initial call.
+    #
+    # @return [Enumerator]
+    #
+    def build_enumerator(cursor:)
+    end
+  end
+end

data/lib/online_migrations/lock_retrier.rb

@@ -91,10 +91,13 @@ module OnlineMigrations
         else
           yield
         end
-      rescue ActiveRecord::LockWaitTimeout
+      rescue ActiveRecord::LockWaitTimeout, ActiveRecord::Deadlocked => e
         if current_attempt <= attempts
           current_delay = delay(current_attempt)
-          Utils.say("Lock timeout. Retrying in #{current_delay} seconds...")
+
+          problem = e.is_a?(ActiveRecord::Deadlocked) ? "Deadlock detected." : "Lock timeout."
+          Utils.say("#{problem} Retrying in #{current_delay} seconds...")
+
           sleep(current_delay)
           retry
         end

data/lib/online_migrations/schema_statements.rb

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

data/lib/online_migrations/shard_aware.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module ShardAware
+    extend ActiveSupport::Concern
+
+    included do
+      before_validation :set_connection_class_name
+    end
+
+    def on_shard_if_present(&block)
+      if shard
+        connection_class.connected_to(shard: shard.to_sym, role: :writing, &block)
+      else
+        yield
+      end
+    end
+
+    def connection_class_name=(value)
+      if value && (klass = value.safe_constantize)
+        if !(klass <= ActiveRecord::Base)
+          raise ArgumentError, "connection_class_name is not an ActiveRecord::Base child class"
+        end
+
+        connection_class = Utils.find_connection_class(klass)
+        super(connection_class.name)
+      end
+    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 set_connection_class_name
+        self.connection_class_name ||= "ActiveRecord::Base"
+      end
+  end
+end

data/lib/online_migrations/version.rb

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

data/lib/online_migrations.rb

@@ -6,7 +6,8 @@ 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/background_data_migrations/migration_helpers"
+require "online_migrations/active_record_batch_enumerator"
 require "online_migrations/schema_statements"
 require "online_migrations/schema_cache"
 require "online_migrations/migration"
@@ -29,7 +30,8 @@ module OnlineMigrations
   autoload :ForeignKeysCollector
   autoload :IndexDefinition
   autoload :CommandChecker
-  autoload :BackgroundMigration
+  autoload :DataMigration
+  autoload :ShardAware
 
   autoload_at "online_migrations/lock_retrier" do
     autoload :LockRetrier
@@ -40,24 +42,21 @@ module OnlineMigrations
 
   autoload :CopyTrigger
 
-  module BackgroundMigrations
+  module BackgroundDataMigrations
     extend ActiveSupport::Autoload
 
     autoload :Config
     autoload :MigrationStatusValidator
-    autoload :MigrationJobStatusValidator
-    autoload :BackgroundMigrationClassValidator
     autoload :BackfillColumn
     autoload :CopyColumn
     autoload :DeleteAssociatedRecords
     autoload :DeleteOrphanedRecords
     autoload :PerformActionOnRelation
     autoload :ResetCounters
-    autoload :MigrationJob
     autoload :Migration
-    autoload :MigrationJobRunner
-    autoload :MigrationRunner
+    autoload :MigrationJob
     autoload :Scheduler
+    autoload :Ticker
   end
 
   module BackgroundSchemaMigrations
@@ -70,6 +69,10 @@ module OnlineMigrations
     autoload :Scheduler
   end
 
+  # Make aliases for less typing.
+  DataMigrations = BackgroundDataMigrations
+  SchemaMigrations = BackgroundSchemaMigrations
+
   class << self
     # @private
     attr_accessor :current_migration
@@ -87,10 +90,10 @@ module OnlineMigrations
     # @option options [String, Symbol, nil] :shard The name of the shard to run
     #   background data migrations on. By default runs on all shards.
     #
-    def run_background_migrations(**options)
-      BackgroundMigrations::Scheduler.run(**options)
+    def run_background_data_migrations(**options)
+      BackgroundDataMigrations::Scheduler.run(**options)
     end
-    alias run_background_data_migrations run_background_migrations
+    alias run_background_migrations run_background_data_migrations
 
     # Run background schema migrations
     #
@@ -122,6 +125,10 @@ module OnlineMigrations
       else
         ActiveRecord::ConnectionAdapters::SchemaCache.prepend(OnlineMigrations::SchemaCache)
       end
+
+      if !ActiveRecord::Batches::BatchEnumerator.method_defined?(:use_ranges)
+        ActiveRecord::Batches::BatchEnumerator.include(OnlineMigrations::ActiveRecordBatchEnumerator)
+      end
     end
   end
 end