online_migrations 0.26.0 → 0.27.0

data/lib/online_migrations/background_data_migrations/migration.rb

@@ -0,0 +1,312 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundDataMigrations
+    # Class representing background data migration.
+    #
+    # @note The records of this class should not be created manually, but via
+    #   `enqueue_background_data_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.
+        "pausing",     # The migration has been told to pause but is finishing work.
+        "paused",      # The migration was paused in the middle of the run by the user.
+        "failed",      # The migration raises an exception when running.
+        "succeeded",   # The migration finished without error.
+        "cancelling",  # The migration has been told to cancel but is finishing work.
+        "cancelled",   # The migration was cancelled by the user.
+      ]
+
+      COMPLETED_STATUSES = ["succeeded", "failed", "cancelled"]
+
+      ACTIVE_STATUSES = [
+        "enqueued",
+        "running",
+        "pausing",
+        "paused",
+        "cancelling",
+      ]
+
+      STOPPING_STATUSES = ["pausing", "cancelling", "cancelled"]
+
+      self.table_name = :background_data_migrations
+      self.ignored_columns += ["parent_id", "batch_column_name", "min_value", "max_value", "rows_count",
+                               "batch_size", "sub_batch_size", "batch_pause", "sub_batch_pause_ms", "composite"]
+
+      scope :queue_order, -> { order(created_at: :asc) }
+      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
+
+      alias_attribute :name, :migration_name
+
+      enum :status, STATUSES.index_with(&:to_s)
+
+      validates :migration_name, presence: true
+      validates :arguments, uniqueness: { scope: [:migration_name, :shard] }
+
+      validates_with MigrationStatusValidator, on: :update
+
+      before_save :set_defaults
+      after_save :instrument_status_change, if: :status_previously_changed?
+
+      # @private
+      def self.normalize_migration_name(migration_name)
+        namespace = ::OnlineMigrations.config.background_data_migrations.migrations_module
+        migration_name.sub(/^(::)?#{namespace}::/, "")
+      end
+
+      def migration_name=(class_name)
+        class_name = class_name.name if class_name.is_a?(Class)
+        write_attribute(:migration_name, self.class.normalize_migration_name(class_name))
+      end
+      alias name= migration_name=
+
+      # Returns whether the migration has been started, which is indicated by the
+      # started_at timestamp being present.
+      #
+      # @return [Boolean] whether the migration was started.
+      #
+      def started?
+        started_at.present?
+      end
+
+      # 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?
+        COMPLETED_STATUSES.include?(status)
+      end
+
+      # Returns whether the migration is active, which is defined as
+      # having a status of enqueued, running, pausing, paused, or cancelling.
+      #
+      # @return [Boolean] whether the migration is active.
+      #
+      def active?
+        ACTIVE_STATUSES.include?(status)
+      end
+
+      # Returns whether the migration is stopping, which is defined as having a status
+      # of pausing or cancelling. The status of cancelled is also considered
+      # stopping since a migration can be cancelled while its job still exists in the
+      # queue, and we want to handle it the same way as a cancelling run.
+      #
+      # @return [Boolean] whether the migration is stopping.
+      #
+      def stopping?
+        STOPPING_STATUSES.include?(status)
+      end
+
+      # Returns whether a migration is stuck, which is defined as having a status of
+      # cancelling or pausing, and not having been updated in the last 5 minutes.
+      #
+      # @return [Boolean] whether the migration is stuck.
+      #
+      def stuck?
+        stuck_timeout = OnlineMigrations.config.background_data_migrations.stuck_timeout
+        (cancelling? || pausing?) && updated_at <= stuck_timeout.ago
+      end
+
+      # @private
+      def start
+        if running? && !started?
+          update!(started_at: Time.current)
+          data_migration.after_start
+          true
+        elsif enqueued?
+          update!(status: :running, started_at: Time.current)
+          data_migration.after_start
+          true
+        else
+          false
+        end
+      end
+
+      # Cancel this data migration. No-op if migration is completed.
+      #
+      # @return [Boolean] whether this data migration was cancelled.
+      #
+      def cancel
+        return false if completed?
+
+        if paused? || stuck?
+          update!(status: :cancelled, finished_at: Time.current)
+        elsif enqueued?
+          cancelled!
+        else
+          cancelling!
+        end
+
+        true
+      end
+
+      # Pause this data migration. No-op if migration is completed.
+      #
+      # @return [Boolean] whether this data migration was paused.
+      #
+      def pause
+        return false if completed?
+
+        if enqueued? || stuck?
+          paused!
+        else
+          pausing!
+        end
+
+        true
+      end
+
+      # Resume this data migration. No-op if migration is not paused.
+      #
+      # @return [Boolean] whether this data migration was resumed.
+      #
+      def resume
+        if paused?
+          enqueued!
+          true
+        else
+          false
+        end
+      end
+
+      # @private
+      def stop
+        return false if completed?
+
+        if cancelling?
+          update!(status: :cancelled, finished_at: Time.current)
+          data_migration.after_cancel
+          data_migration.after_complete
+        elsif pausing?
+          paused!
+          data_migration.after_pause
+        end
+
+        data_migration.after_stop
+        true
+      end
+
+      # @private
+      def complete
+        return false if completed?
+
+        if running?
+          update!(status: :succeeded, finished_at: Time.current)
+          data_migration.after_complete
+        elsif pausing?
+          update!(status: :paused, finished_at: Time.current)
+        elsif cancelling?
+          update!(status: :cancelled, finished_at: Time.current)
+          data_migration.after_complete
+        end
+
+        true
+      end
+
+      # @private
+      def persist_progress(cursor, number_of_ticks, duration)
+        update!(
+          cursor: cursor,
+          tick_count: tick_count + number_of_ticks,
+          time_running: time_running + duration
+        )
+      end
+
+      # @private
+      def persist_error(error)
+        backtrace = error.backtrace
+        backtrace_cleaner = OnlineMigrations.config.backtrace_cleaner
+        backtrace = backtrace_cleaner.clean(backtrace) if backtrace_cleaner
+
+        update!(
+          status: :failed,
+          finished_at: Time.current,
+          error_class: error.class.name,
+          error_message: error.message,
+          backtrace: backtrace
+        )
+      end
+
+      # Returns whether this migration is pausable.
+      #
+      def pausable?
+        true
+      end
+
+      # Returns the progress of the data migration.
+      #
+      # @return [Float, nil]
+      #   - when background migration is configured to not track progress, returns `nil`
+      #   - otherwise returns value in range from 0.0 to 100.0
+      #
+      def progress
+        if succeeded?
+          100.0
+        elsif enqueued? || tick_total == 0
+          0.0
+        elsif tick_total
+          ([tick_count.to_f / tick_total, 1.0].min * 100)
+        end
+      end
+
+      # Returns data migration associated with this migration.
+      #
+      # @return [OnlineMigrations::DataMigration]
+      #
+      def data_migration
+        @data_migration ||= begin
+          klass = DataMigration.named(migration_name)
+          klass.new(*arguments)
+        end
+      end
+
+      # Mark this migration as ready to be processed again.
+      #
+      # This method marks failed migrations as ready to be processed again, and
+      # they will be picked up on the next Scheduler run.
+      #
+      def retry
+        if failed?
+          update!(
+            status: :enqueued,
+            started_at: nil,
+            finished_at: nil,
+            error_class: nil,
+            error_message: nil,
+            backtrace: nil,
+            jid: nil
+          )
+          true
+        else
+          false
+        end
+      end
+
+      private
+        def set_defaults
+          config = ::OnlineMigrations.config.background_data_migrations
+          self.max_attempts ||= config.max_attempts
+          self.tick_total ||= on_shard_if_present do
+            data_migration.count
+          end
+        end
+
+        def instrument_status_change
+          payload = { migration: self }
+          if running?
+            ActiveSupport::Notifications.instrument("started.background_data_migrations", payload)
+          elsif succeeded?
+            ActiveSupport::Notifications.instrument("completed.background_data_migrations", payload)
+          end
+        end
+    end
+  end
+end

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

@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 
 module OnlineMigrations
-  module BackgroundMigrations
+  module BackgroundDataMigrations
     module MigrationHelpers
       # Backfills column data using background migrations.
       #
@@ -13,7 +13,7 @@ module OnlineMigrations
       # @param options [Hash] used to control the behavior of background migration.
       #     See `#enqueue_background_data_migration`
       #
-      # @return [OnlineMigrations::BackgroundMigrations::Migration]
+      # @return [OnlineMigrations::BackgroundDataMigrations::Migration]
       #
       # @example
       #   backfill_column_in_background(:users, :admin, false)
@@ -69,7 +69,7 @@ module OnlineMigrations
       # @param options [Hash] used to control the behavior of background migration.
       #     See `#enqueue_background_data_migration`
       #
-      # @return [OnlineMigrations::BackgroundMigrations::Migration]
+      # @return [OnlineMigrations::BackgroundDataMigrations::Migration]
       #
       # @example
       #   backfill_column_for_type_change_in_background(:files, :size)
@@ -108,7 +108,14 @@ module OnlineMigrations
         end
 
         tmp_columns = column_names.map { |column_name| "#{column_name}_for_type_change" }
-        model_name = model_name.name if model_name.is_a?(Class)
+
+        if model_name
+          model_name = model_name.name if model_name.is_a?(Class)
+          connection_class_name = Utils.find_connection_class(model_name.constantize).name
+        end
+
+        # model_name = model_name.name if model_name.is_a?(Class)
+        # connection_class = Utils.find_connection_class(model_name.constantize) if model_name
 
         enqueue_background_data_migration(
           "CopyColumn",
@@ -117,6 +124,7 @@ module OnlineMigrations
           tmp_columns,
           model_name,
           type_cast_functions,
+          connection_class_name: connection_class_name,
           **options
         )
       end
@@ -134,7 +142,7 @@ module OnlineMigrations
       # @param options [Hash] used to control the behavior of background migration.
       #     See `#enqueue_background_data_migration`
       #
-      # @return [OnlineMigrations::BackgroundMigrations::Migration]
+      # @return [OnlineMigrations::BackgroundDataMigrations::Migration]
       #
       # @example
       #   copy_column_in_background(:users, :id, :id_for_type_change)
@@ -189,7 +197,7 @@ module OnlineMigrations
       # @param options [Hash] used to control the behavior of background migration.
       #     See `#enqueue_background_data_migration`
       #
-      # @return [OnlineMigrations::BackgroundMigrations::Migration]
+      # @return [OnlineMigrations::BackgroundDataMigrations::Migration]
       #
       # @example
       #     reset_counters_in_background("User", :projects, :friends, touch: true)
@@ -226,7 +234,7 @@ module OnlineMigrations
       # @param options [Hash] used to control the behavior of background migration.
       #     See `#enqueue_background_data_migration`
       #
-      # @return [OnlineMigrations::BackgroundMigrations::Migration]
+      # @return [OnlineMigrations::BackgroundDataMigrations::Migration]
       #
       # @example
       #     delete_orphaned_records_in_background("Post", :author)
@@ -255,7 +263,7 @@ module OnlineMigrations
       # @param options [Hash] used to control the behavior of background migration.
       #     See `#enqueue_background_data_migration`
       #
-      # @return [OnlineMigrations::BackgroundMigrations::Migration]
+      # @return [OnlineMigrations::BackgroundDataMigrations::Migration]
       #
       # @example
       #     delete_associated_records_in_background("Link", 1, :clicks)
@@ -286,7 +294,7 @@ module OnlineMigrations
       # @param options [Hash] used to control the behavior of background migration.
       #     See `#enqueue_background_data_migration`
       #
-      # @return [OnlineMigrations::BackgroundMigrations::Migration]
+      # @return [OnlineMigrations::BackgroundDataMigrations::Migration]
       #
       # @example Delete records
       #     perform_action_on_relation_in_background("User", { banned: true }, :delete_all)
@@ -328,56 +336,64 @@ module OnlineMigrations
       # based on the current migration settings and the previous batch bounds. Each job's execution status
       # is tracked in the database as the migration runs.
       #
-      # @param migration_name [String, Class] Background migration job class name
-      # @param arguments [Array] Extra arguments to pass to the job instance when the migration runs
-      # @option options [Symbol, String] :batch_column_name (primary key) Column name the migration will batch over
-      # @option options [Integer] :min_value Value in the column the batching will begin at,
-      #     defaults to `SELECT MIN(batch_column_name)`
-      # @option options [Integer] :max_value Value in the column the batching will end at,
-      #     defaults to `SELECT MAX(batch_column_name)`
-      # @option options [Integer] :batch_size (1_000) Number of rows to process in a single background migration run
-      # @option options [Integer] :sub_batch_size (100) Smaller batches size that the batches will be divided into
-      # @option options [Integer] :batch_pause (0) Pause interval between each background migration job's execution (in seconds)
-      # @option options [Integer] :sub_batch_pause_ms (100) Number of milliseconds to sleep between each sub_batch execution
-      # @option options [Integer] :batch_max_attempts (5) Maximum number of batch run attempts
-      #
-      # @return [OnlineMigrations::BackgroundMigrations::Migration]
+      # @param migration_name [String, Class] Background migration class name
+      # @param arguments [Array] Extra arguments to pass to the migration instance when the migration runs
+      # @option options [Integer] :max_attempts (5) Maximum number of batch run attempts
+      # @option options [String, nil] :connection_class_name Class name to use to get connections
+      #
+      # @return [OnlineMigrations::BackgroundDataMigrations::Migration]
       #
       # @example
-      #   enqueue_background_data_migration("BackfillProjectIssuesCount",
-      #       batch_size: 10_000, batch_max_attempts: 10)
+      #   enqueue_background_data_migration("BackfillProjectIssuesCount")
       #
       #   # Given the background migration exists:
       #
-      #   class BackfillProjectIssuesCount < OnlineMigrations::BackgroundMigration
-      #     def relation
-      #       Project.all
+      #   class BackfillProjectIssuesCount < OnlineMigrations::DataMigration
+      #     def collection
+      #       Project.in_batches(of: 100)
       #     end
       #
-      #     def process_batch(projects)
+      #     def process(projects)
       #       projects.update_all(
       #         "issues_count = (SELECT COUNT(*) FROM issues WHERE issues.project_id = projects.id)"
       #       )
       #     end
       #
-      #     # To be able to track progress, you need to define this method
+      #     # To be able to track progress, you need to define this method.
       #     def count
       #       Project.maximum(:id)
       #     end
       #   end
       #
-      # @note For convenience, the enqueued background migration is run inline
+      # @note For convenience, the enqueued background data migration is run inline
       #     in development and test environments
       #
       def enqueue_background_data_migration(migration_name, *arguments, **options)
-        migration = create_background_data_migration(migration_name, *arguments, **options)
+        options.assert_valid_keys(:max_attempts, :connection_class_name)
+
+        migration_name = migration_name.name if migration_name.is_a?(Class)
+        options[:connection_class_name] ||= compute_connection_class_name(migration_name, arguments)
 
-        if Utils.run_background_migrations_inline?
-          runner = MigrationRunner.new(migration)
-          runner.run_all_migration_jobs
+        if Utils.multiple_databases? && !options[:connection_class_name]
+          raise ArgumentError, "You must pass a :connection_class_name when using multiple databases."
         end
 
-        migration
+        connection_class = options[:connection_class_name].constantize
+        shards = Utils.shard_names(connection_class)
+        shards = [nil] if shards.size == 1
+
+        shards.each do |shard|
+          # Can't use `find_or_create_by` or hash syntax here, because it does not correctly work with json `arguments`.
+          migration = Migration.where(migration_name: migration_name, shard: shard).where("arguments = ?", arguments.to_json).first
+          migration ||= Migration.create!(**options, migration_name: migration_name, arguments: arguments, shard: shard)
+
+          if Utils.run_background_migrations_inline?
+            job = OnlineMigrations.config.background_data_migrations.job
+            job.constantize.perform_inline(migration.id)
+          end
+        end
+
+        true
       end
       alias enqueue_background_migration enqueue_background_data_migration
 
@@ -418,45 +434,40 @@ module OnlineMigrations
 
         if arguments
           arguments = Array(arguments)
-          migration = Migration.parents.for_configuration(migration_name, arguments).first
+          migrations = Migration.for_configuration(migration_name, arguments).to_a
           configuration[:arguments] = arguments.to_json
         else
-          migration = Migration.parents.for_migration_name(migration_name).first
+          migrations = Migration.for_migration_name(migration_name).to_a
         end
 
-        if migration.nil?
-          Utils.raise_in_prod_or_say_in_dev("Could not find background data migration for the given configuration: #{configuration}")
-        elsif !migration.succeeded?
-          raise "Expected background data migration for the given configuration to be marked as 'succeeded', " \
-                "but it is '#{migration.status}': #{configuration}"
+        if migrations.empty?
+          Utils.raise_in_prod_or_say_in_dev("Could not find background data migration(s) for the given configuration: #{configuration}.")
+        elsif !migrations.all?(&:succeeded?)
+          raise "Expected background data migration(s) for the given configuration to be marked as 'succeeded': #{configuration}."
         end
       end
       alias ensure_background_migration_succeeded ensure_background_data_migration_succeeded
 
-      # @private
-      def create_background_data_migration(migration_name, *arguments, **options)
-        options.assert_valid_keys(:batch_column_name, :min_value, :max_value, :batch_size, :sub_batch_size,
-            :batch_pause, :sub_batch_pause_ms, :batch_max_attempts)
+      private
+        def compute_connection_class_name(migration_name, arguments)
+          klass = DataMigration.named(migration_name)
+          data_migration = klass.new(*arguments)
 
-        migration_name = migration_name.name if migration_name.is_a?(Class)
+          collection = data_migration.collection
 
-        # Can't use `find_or_create_by` or hash syntax here, because it does not correctly work with json `arguments`.
-        existing_migration = Migration.find_by("migration_name = ? AND arguments = ? AND shard IS NULL", migration_name, arguments.to_json)
-        return existing_migration if existing_migration
-
-        Migration.create!(**options, migration_name: migration_name, arguments: arguments, shard: nil) do |migration|
-          shards = Utils.shard_names(migration.migration_model)
-          if shards.size > 1
-            migration.children = shards.map do |shard|
-              child = migration.dup
-              child.shard = shard
-              child
+          model =
+            case collection
+            when ActiveRecord::Relation then collection.model
+            when ActiveRecord::Batches::BatchEnumerator then collection.relation.model
             end
 
-            migration.composite = true
+          if model
+            connection_class = Utils.find_connection_class(model)
+            connection_class.name
           end
+        rescue NotImplementedError
+          nil
         end
-      end
     end
   end
 end