online_migrations 0.26.0 → 0.27.0

data/lib/online_migrations/background_data_migrations/migration_job.rb

@@ -0,0 +1,158 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundDataMigrations
+    # Sidekiq job responsible for running background data migrations.
+    class MigrationJob
+      include Sidekiq::IterableJob
+
+      sidekiq_retry_in do |count, _exception, jobhash|
+        migration_id = jobhash["args"].fetch(0)
+        migration = Migration.find(migration_id)
+
+        if count + 1 >= migration.max_attempts
+          :kill
+        end
+      end
+
+      sidekiq_retries_exhausted do |jobhash, exception|
+        migration_id = jobhash["args"].fetch(0)
+        migration = Migration.find(migration_id)
+        migration.persist_error(exception)
+
+        OnlineMigrations.config.background_data_migrations.error_handler.call(exception, migration)
+      end
+
+      TICKER_INTERVAL = 5 # seconds
+
+      def initialize
+        super
+
+        @migration = nil
+        @data_migration = nil
+
+        @ticker = Ticker.new(TICKER_INTERVAL) do |ticks, duration|
+          # TODO: use 'cursor' accessor from sidekiq in the future.
+          # https://github.com/sidekiq/sidekiq/pull/6606
+          @migration.persist_progress(@_cursor, ticks, duration)
+          @migration.reload
+        end
+
+        @throttle_checked_at = current_time
+      end
+
+      def on_start
+        @migration.start
+      end
+
+      def around_iteration(&block)
+        @migration.on_shard_if_present(&block)
+      end
+
+      def on_resume
+        @data_migration.after_resume
+      end
+
+      def on_stop
+        @ticker.persist
+        @migration.stop
+      end
+
+      def on_complete
+        # Job was manually cancelled.
+        @migration.cancel if cancelled?
+
+        @migration.complete
+      end
+
+      def build_enumerator(migration_id, cursor:)
+        @migration = BackgroundDataMigrations::Migration.find(migration_id)
+        cursor ||= @migration.cursor
+
+        @migration.on_shard_if_present do
+          @data_migration = @migration.data_migration
+          collection_enum = @data_migration.build_enumerator(cursor: cursor)
+
+          if collection_enum
+            if !collection_enum.is_a?(Enumerator)
+              raise ArgumentError, <<~MSG.squish
+                #{@data_migration.class.name}#build_enumerator must return an Enumerator,
+                got #{collection_enum.class.name}.
+              MSG
+            end
+
+            collection_enum
+          else
+            collection = @data_migration.collection
+
+            case collection
+            when ActiveRecord::Relation
+              options = {
+                cursor: cursor,
+                batch_size: @data_migration.class.active_record_enumerator_batch_size || 100,
+              }
+              active_record_records_enumerator(collection, **options)
+            when ActiveRecord::Batches::BatchEnumerator
+              if collection.start || collection.finish
+                raise ArgumentError, <<~MSG.squish
+                  #{@data_migration.class.name}#collection does not support
+                  a batch enumerator with the "start" or "finish" options.
+                MSG
+              end
+
+              active_record_relations_enumerator(
+                collection.relation,
+                batch_size: collection.batch_size,
+                cursor: cursor,
+                use_ranges: collection.use_ranges
+              )
+            when Array
+              array_enumerator(collection, cursor: cursor)
+            else
+              raise ArgumentError, <<~MSG.squish
+                #{@data_migration.class.name}#collection must be either an ActiveRecord::Relation,
+                ActiveRecord::Batches::BatchEnumerator, or Array.
+              MSG
+            end
+          end
+        end
+      end
+
+      def each_iteration(item, _migration_id)
+        if @migration.cancelling? || @migration.pausing? || @migration.paused?
+          # Finish this exact sidekiq job. When the migration is paused
+          # and will be resumed, a new job will be enqueued.
+          finished = true
+          throw :abort, finished
+        elsif should_throttle?
+          ActiveSupport::Notifications.instrument("throttled.background_data_migrations", migration: @migration)
+          finished = false
+          throw :abort, finished
+        else
+          @data_migration.around_process do
+            @migration.data_migration.process(item)
+
+            pause = OnlineMigrations.config.background_data_migrations.iteration_pause
+            sleep(pause) if pause > 0
+          end
+          @ticker.tick
+        end
+      end
+
+      private
+        THROTTLE_CHECK_INTERVAL = 5 # seconds
+        private_constant :THROTTLE_CHECK_INTERVAL
+
+        def should_throttle?
+          if current_time - @throttle_checked_at >= THROTTLE_CHECK_INTERVAL
+            @throttle_checked_at = current_time
+            OnlineMigrations.config.throttler.call
+          end
+        end
+
+        def current_time
+          ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
+        end
+    end
+  end
+end

data/lib/online_migrations/background_data_migrations/migration_status_validator.rb

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundDataMigrations
+    # @private
+    class MigrationStatusValidator < ActiveModel::Validator
+      # Valid status transitions a Migration can make.
+      VALID_STATUS_TRANSITIONS = {
+        # enqueued -> running occurs when the migration starts performing.
+        # enqueued -> paused occurs when the migration is paused before starting.
+        # enqueued -> cancelled occurs when the migration is cancelled before starting.
+        # enqueued -> failed occurs when the migration job fails to be enqueued, or
+        #   if the migration is deleted before is starts running.
+        "enqueued" => ["running", "paused", "cancelled", "failed"],
+        # running -> succeeded occurs when the migration completes successfully.
+        # running -> pausing occurs when a user pauses the migration as it's performing.
+        # running -> cancelling occurs when a user cancels the migration as it's performing.
+        # running -> failed occurs when the job raises an exception when running.
+        "running" => [
+          "succeeded",
+          "pausing",
+          "cancelling",
+          "failed",
+        ],
+        # pausing -> paused occurs when the migration actually halts performing and
+        #   occupies a status of paused.
+        # pausing -> cancelling occurs when the user cancels a migration immediately
+        #   after it was paused, such that the migration had not actually halted yet.
+        # pausing -> succeeded occurs when the migration completes immediately after
+        #   being paused. This can happen if the migration is on its last iteration
+        #   when it is paused, or if the migration is paused after enqueue but has
+        #   nothing in its collection to process.
+        # pausing -> failed occurs when the job raises an exception after the
+        #   user has paused it.
+        "pausing" => ["paused", "cancelling", "succeeded", "failed"],
+        # paused -> enqueued occurs when the migration is resumed after being paused.
+        # paused -> cancelled when the user cancels the migration after it is paused.
+        "paused" => ["enqueued", "cancelled"],
+        # failed -> enqueued occurs when the migration is retried after encounting an error.
+        "failed" => ["enqueued"],
+        # cancelling -> cancelled occurs when the migration actually halts performing
+        #   and occupies a status of cancelled.
+        # cancelling -> succeeded occurs when the migration completes immediately after
+        #   being cancelled. See description for pausing -> succeeded.
+        # cancelling -> failed occurs when the job raises an exception after the
+        #   user has cancelled it.
+        "cancelling" => ["cancelled", "succeeded", "failed"],
+      }
+
+      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 data migration from status '#{previous_status}' to '#{new_status}'"
+          )
+        end
+      end
+    end
+  end
+end

data/lib/online_migrations/{background_migrations → background_data_migrations}/perform_action_on_relation.rb

@@ -1,9 +1,9 @@
 # frozen_string_literal: true
 
 module OnlineMigrations
-  module BackgroundMigrations
+  module BackgroundDataMigrations
     # @private
-    class PerformActionOnRelation < BackgroundMigration
+    class PerformActionOnRelation < DataMigration
       attr_reader :model, :conditions, :action, :options
 
       def initialize(model_name, conditions, action, options = {})
@@ -13,11 +13,11 @@ module OnlineMigrations
         @options = options.symbolize_keys
       end
 
-      def relation
-        model.unscoped.where(conditions)
+      def collection
+        model.unscoped.where(conditions).in_batches(of: 100)
       end
 
-      def process_batch(relation)
+      def process(relation)
         case action
         when :update_all
           updates = options.fetch(:updates)

data/lib/online_migrations/{background_migrations → background_data_migrations}/reset_counters.rb

@@ -1,9 +1,9 @@
 # frozen_string_literal: true
 
 module OnlineMigrations
-  module BackgroundMigrations
+  module BackgroundDataMigrations
     # @private
-    class ResetCounters < BackgroundMigration
+    class ResetCounters < DataMigration
       attr_reader :model, :counters, :touch
 
       def initialize(model_name, counters, options = {})
@@ -12,11 +12,11 @@ module OnlineMigrations
         @touch = options[:touch]
       end
 
-      def relation
-        model.unscoped
+      def collection
+        model.unscoped.in_batches(of: 100)
       end
 
-      def process_batch(relation)
+      def process(relation)
         updates = counters.map do |counter_association|
           has_many_association = has_many_association(counter_association)
 

data/lib/online_migrations/background_data_migrations/scheduler.rb

@@ -0,0 +1,78 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundDataMigrations
+    # Class responsible for scheduling background data migrations.
+    #
+    # 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_data_migrations"
+    #   end
+    #
+    # @example Specific shard
+    #   every 1.minute do
+    #     runner "OnlineMigrations.run_background_data_migrations(shard: :shard_two)"
+    #   end
+    #
+    # @example Custom concurrency
+    #   every 1.minute do
+    #     # Allow to run 2 data migrations in parallel.
+    #     runner "OnlineMigrations.run_background_data_migrations(concurrency: 2)"
+    #   end
+    #
+    class Scheduler
+      def self.run(**options)
+        new.run(**options)
+      end
+
+      # Runs Scheduler
+      def run(shard: nil, concurrency: 1)
+        relation = Migration.queue_order
+        relation = relation.where(shard: shard) if shard
+
+        with_lock do
+          running = relation.running
+          enqueued = relation.enqueued
+
+          # Ensure no more than 'concurrency' migrations are running at the same time.
+          remaining_to_enqueue = concurrency - running.count
+          if remaining_to_enqueue > 0
+            migrations_to_enqueue = enqueued.limit(remaining_to_enqueue)
+            migrations_to_enqueue.each do |migration|
+              enqueue_migration(migration)
+            end
+          end
+        end
+
+        true
+      end
+
+      private
+        def with_lock(&block)
+          # Don't lock the whole table if we can lock only a single record (must be always the same).
+          first_record = Migration.queue_order.first
+          if first_record
+            first_record.with_lock(&block)
+          else
+            Migration.transaction do
+              Migration.connection.execute("LOCK #{Migration.table_name} IN ACCESS EXCLUSIVE MODE")
+              yield
+            end
+          end
+        end
+
+        def enqueue_migration(migration)
+          job = OnlineMigrations.config.background_data_migrations.job
+          job_class = job.constantize
+
+          jid = job_class.perform_async(migration.id)
+          if jid
+            migration.update!(status: :running, jid: jid)
+          end
+        end
+    end
+  end
+end

data/lib/online_migrations/background_data_migrations/ticker.rb

@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundDataMigrations
+    # This class encapsulates the logic behind updating the tick counter.
+    #
+    # It's initialized with a duration for the throttle, and a block to persist
+    # the number of ticks to increment.
+    #
+    # When +tick+ is called, the block will be called with the increment,
+    # provided the duration since the last update (or initialization) has been
+    # long enough.
+    #
+    # To not lose any increments, +persist+ should be used, which may call the
+    # block with any leftover ticks.
+    #
+    # @private
+    class Ticker
+      # Creates a Ticker that will call the block each time +tick+ is called,
+      # unless the tick is being throttled.
+      #
+      # @param interval [ActiveSupport::Duration, Numeric] Duration
+      #   since initialization or last call that will cause a throttle.
+      # @yieldparam ticks [Integer] the increment in ticks to be persisted.
+      #
+      def initialize(interval, &block)
+        @interval = interval
+        @block = block
+        @last_persisted_at = Time.current
+        @ticks_recorded = 0
+      end
+
+      # Increments the tick count by one, and may persist the new value if the
+      # threshold duration has passed since initialization or the tick count was
+      # last persisted.
+      #
+      def tick
+        @ticks_recorded += 1
+        persist if persist?
+      end
+
+      # Persists the tick increments by calling the block passed to the
+      # initializer. This is idempotent in the sense that calling it twice in a
+      # row will call the block at most once (if it had been throttled).
+      #
+      def persist
+        return if @ticks_recorded == 0
+
+        now = Time.current
+        duration = now - @last_persisted_at
+        @last_persisted_at = now
+        @block.call(@ticks_recorded, duration)
+        @ticks_recorded = 0
+      end
+
+      private
+        def persist?
+          Time.now - @last_persisted_at >= @interval
+        end
+    end
+  end
+end

data/lib/online_migrations/background_schema_migrations/config.rb

@@ -4,9 +4,9 @@ module OnlineMigrations
   module BackgroundSchemaMigrations
     # Class representing configuration options for background schema migrations.
     class Config
-      # Maximum number of run attempts
+      # Maximum number of run attempts.
       #
-      # When attempts are exhausted, the migration is marked as failed.
+      # When attempts are exhausted, the schema migration is marked as failed.
       # @return [Integer] defaults to 5
       #
       attr_accessor :max_attempts

data/lib/online_migrations/background_schema_migrations/migration.rb

@@ -8,13 +8,15 @@ module OnlineMigrations
     #   `enqueue_background_schema_migration` helper inside migrations.
     #
     class Migration < ApplicationRecord
+      include ShardAware
+
       STATUSES = [
-        :enqueued,    # The migration has been enqueued by the user.
-        :running,     # The migration is being performed by a migration executor.
-        :errored,     # The migration raised an error during last run.
-        :failed,      # The migration raises an error when running and retry attempts exceeded.
-        :succeeded,   # The migration finished without error.
-        :cancelled,   # The migration was cancelled by the user.
+        "enqueued",    # The migration has been enqueued by the user.
+        "running",     # The migration is being performed by a migration executor.
+        "errored",     # The migration raised an error during last run.
+        "failed",      # The migration raises an error when running and retry attempts exceeded.
+        "succeeded",   # The migration finished without error.
+        "cancelled",   # The migration was cancelled by the user.
       ]
 
       MAX_IDENTIFIER_LENGTH = 63
@@ -22,36 +24,12 @@ module OnlineMigrations
       self.table_name = :background_schema_migrations
 
       scope :queue_order, -> { order(created_at: :asc) }
-      scope :parents, -> { where(parent_id: nil) }
-      scope :runnable, -> { where(composite: false) }
       scope :active, -> { where(status: [:enqueued, :running, :errored]) }
-      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
-        stuck_sql = connection.unprepared_statement { stuck.to_sql }
-
-        from(Arel.sql(<<~SQL))
-          (
-            (SELECT * FROM background_schema_migrations WHERE NOT composite AND status = 'errored')
-            UNION
-            (#{stuck_sql})
-          ) AS #{table_name}
-        SQL
-      end
 
       alias_attribute :name, :migration_name
 
       enum :status, STATUSES.index_with(&:to_s)
 
-      belongs_to :parent, class_name: name, optional: true, inverse_of: :children
-      has_many :children, class_name: name, foreign_key: :parent_id, inverse_of: :parent
-
       validates :table_name, presence: true, length: { maximum: MAX_IDENTIFIER_LENGTH }
       validates :definition, presence: true
       validates :migration_name, presence: true, uniqueness: {
@@ -65,70 +43,53 @@ module OnlineMigrations
         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
 
+      # Returns whether the migration is completed, which is defined as
+      # having a status of succeeded, failed, or cancelled.
+      #
+      # @return [Boolean] whether the migration is completed.
+      #
       def completed?
-        succeeded? || failed?
+        succeeded? || failed? || cancelled?
       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
+      # Returns whether the migration is active, which is defined as
+      # having a status of enqueued, or running.
+      #
+      # @return [Boolean] whether the migration is active.
+      #
+      def active?
+        enqueued? || running?
       end
+
       alias cancel cancelled!
 
+      # Returns whether this migration is pausable.
+      #
       def pausable?
         false
       end
 
-      def can_be_paused?
-        false
-      end
-
-      def can_be_cancelled?
-        !succeeded? && !cancelled?
-      end
-
-      # Returns the progress of the background schema migration.
+      # Dummy method to support the same interface as background data migrations.
       #
-      # @return [Float] value in range from 0.0 to 100.0
+      # @return [nil]
       #
       def progress
-        if succeeded?
-          100.0
-        elsif composite?
-          progresses = children.map(&:progress)
-          # There should not be composite migrations without children,
-          # but children may be deleted for some reason, so we need to
-          # make a check to avoid 0 division error.
-          if progresses.any?
-            (progresses.sum.to_f / progresses.size).round(2)
-          else
-            0.0
-          end
-        else
-          0.0
-        end
       end
 
       # Whether the migration is considered stuck (is running for some configured time).
       #
       def stuck?
-        # Composite migrations are not considered stuck.
-        return false if composite?
-
-        stuck_timeout = (statement_timeout || 1.day) + 10.minutes
-        running? && updated_at <= stuck_timeout.seconds.ago
+        if index_addition?
+          running? && !index_build_in_progress?
+        else
+          stuck_timeout = (statement_timeout || 1.day) + 10.minutes
+          running? && updated_at <= stuck_timeout.seconds.ago
+        end
       end
 
       # Mark this migration as ready to be processed again.
@@ -136,27 +97,16 @@ module OnlineMigrations
       # This is used to manually retrying failed migrations.
       #
       def retry
-        if composite? && failed?
-          transaction do
-            update!(status: :enqueued, finished_at: nil)
-            children.failed.each(&:retry)
-          end
-
-          true
-        elsif failed?
-          transaction do
-            parent.update!(status: :enqueued, finished_at: nil) if parent
-
-            update!(
-              status: :enqueued,
-              attempts: 0,
-              started_at: nil,
-              finished_at: nil,
-              error_class: nil,
-              error_message: nil,
-              backtrace: nil
-            )
-          end
+        if failed?
+          update!(
+            status: :enqueued,
+            attempts: 0,
+            started_at: nil,
+            finished_at: nil,
+            error_class: nil,
+            error_message: nil,
+            backtrace: nil
+          )
 
           true
         else
@@ -168,15 +118,6 @@ module OnlineMigrations
         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
@@ -184,7 +125,7 @@ module OnlineMigrations
 
       # @private
       def run
-        on_shard do
+        on_shard_if_present do
           connection = connection_class.connection
 
           connection.with_lock_retries do
@@ -214,28 +155,11 @@ module OnlineMigrations
       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
+          on_shard_if_present do
             if !connection_class.connection.table_exists?(table_name)
               errors.add(:table_name, "'#{table_name}' does not exist")
             end
@@ -248,9 +172,13 @@ module OnlineMigrations
           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)
+        def index_build_in_progress?
+          indexes_in_progress = connection_class.connection.select_values(<<~SQL)
+            SELECT index_relid::regclass::text
+            FROM pg_stat_progress_create_index
+          SQL
+
+          indexes_in_progress.include?(name)
         end
 
         def with_statement_timeout(connection, timeout)