online_migrations 0.1.0

data/lib/online_migrations/background_migrations/copy_column.rb

@@ -0,0 +1,90 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundMigrations
+    # @private
+    class CopyColumn < BackgroundMigration
+      attr_reader :table_name, :copy_from, :copy_to, :model_name, :type_cast_functions
+
+      def initialize(table_name, copy_from, copy_to, model_name = nil, type_cast_functions = {})
+        @table_name = table_name
+
+        if copy_from.is_a?(Array) && type_cast_functions && !type_cast_functions.is_a?(Hash)
+          raise ArgumentError, "type_cast_functions must be a Hash"
+        end
+
+        @copy_from = Array.wrap(copy_from)
+        @copy_to = Array.wrap(copy_to)
+
+        if @copy_from.size != @copy_to.size
+          raise ArgumentError, "Number of source and destination columns must match"
+        end
+
+        @model_name = model_name
+        @type_cast_functions = type_cast_functions
+      end
+
+      def relation
+        relation = model
+          .where(Hash[copy_to.map { |to_column| [to_column, nil] }])
+
+        Utils.ar_where_not_multiple_conditions(
+          relation,
+          copy_from.map { |from_column| [from_column, nil] }.to_h
+        )
+      end
+
+      def process_batch(relation)
+        arel = relation.arel
+        arel_table = relation.arel_table
+
+        old_values = copy_from.map do |from_column|
+          old_value = arel_table[from_column]
+          if (type_cast_function = type_cast_functions[from_column])
+            if Utils.ar_version <= 5.2
+              # ActiveRecord <= 5.2 does not support quoting of Arel::Nodes::NamedFunction
+              old_value = Arel.sql("#{type_cast_function}(#{connection.quote_column_name(from_column)})")
+            else
+              old_value = Arel::Nodes::NamedFunction.new(type_cast_function, [old_value])
+            end
+          end
+          old_value
+        end
+
+        if Utils.ar_version <= 4.2
+          stmt = Arel::UpdateManager.new(arel.engine)
+        else
+          stmt = Arel::UpdateManager.new
+        end
+
+        stmt.table(arel_table)
+        stmt.wheres = arel.constraints
+
+        updates = copy_to.zip(old_values).map { |to_column, old_value| [arel_table[to_column], old_value] }
+        stmt.set(updates)
+
+        connection.update(stmt)
+      end
+
+      def count
+        # Exact counts are expensive on large tables, since PostgreSQL
+        # needs to do a full scan. An estimated count should give a pretty decent
+        # approximation of rows count in this case.
+        Utils.estimated_count(connection, table_name)
+      end
+
+      private
+        def model
+          @model ||= if model_name.present?
+                       Object.const_get(model_name, false)
+                     else
+                       Utils.define_model(ActiveRecord::Base.connection, table_name)
+                     end
+        end
+
+        def connection
+          model.connection
+        end
+    end
+  end
+end

data/lib/online_migrations/background_migrations/migration.rb

@@ -0,0 +1,210 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundMigrations
+    class Migration < ActiveRecord::Base
+      STATUSES = [
+        :enqueued,    # The migration has been enqueued by the user.
+        :running,     # The migration is being performed by a migration executor.
+        :paused,      # The migration was paused in the middle of the run by the user.
+        :finishing,   # The migration is being manually finishing inline by the user.
+        :failed,      # The migration raises an exception when running.
+        :succeeded,   # The migration finished without error.
+      ]
+
+      self.table_name = :background_migrations
+
+      scope :queue_order, -> { order(created_at: :asc) }
+      scope :active, -> { where(status: [statuses[:enqueued], statuses[:running]]) }
+      scope :for_migration_name, ->(migration_name) { where(migration_name: normalize_migration_name(migration_name)) }
+      scope :for_configuration, ->(migration_name, arguments) do
+        for_migration_name(migration_name).where("arguments = ?", arguments.to_json)
+      end
+
+      enum status: STATUSES.map { |status| [status, status.to_s] }.to_h
+
+      has_many :migration_jobs
+
+      validates :migration_name, :batch_column_name, presence: true
+
+      validates :batch_pause, :min_value, :max_value, :batch_size, :sub_batch_size,
+                  presence: true, numericality: { greater_than: 0 }
+
+      validates :sub_batch_pause_ms, presence: true, numericality: { greater_than_or_equal_to: 0 }
+      validates :rows_count, numericality: { greater_than_or_equal_to: 0 }, allow_nil: true
+      validates :arguments, uniqueness: { scope: :migration_name }
+
+      validate :validate_batch_column_values
+      validate :validate_batch_sizes
+      validate :validate_jobs_status, if: :status_changed?
+
+      validates_with BackgroundMigrationClassValidator
+      validates_with MigrationStatusValidator, on: :update
+
+      before_validation :set_defaults
+
+      # @private
+      def self.normalize_migration_name(migration_name)
+        namespace = ::OnlineMigrations.config.background_migrations.migrations_module
+        migration_name.sub(/^(::)?#{namespace}::/, "")
+      end
+
+      def migration_name=(class_name)
+        write_attribute(:migration_name, self.class.normalize_migration_name(class_name))
+      end
+
+      def completed?
+        succeeded? || failed?
+      end
+
+      def last_job
+        migration_jobs.order(max_value: :desc).first
+      end
+
+      def last_completed_job
+        migration_jobs.completed.order(finished_at: :desc).first
+      end
+
+      # Returns the progress of the background migration.
+      #
+      # @return [Float, nil]
+      #   - when background migration is configured to not to track progress, returns `nil`
+      #   - otherwise returns value in range of 0.0 and 1.0
+      #
+      def progress
+        if succeeded?
+          1.0
+        elsif rows_count
+          jobs_rows_count = migration_jobs.succeeded.sum(:batch_size)
+          # The last migration job may need to process the amount of rows
+          # less than the batch size, so we can get a value > 1.0.
+          [jobs_rows_count.to_f / rows_count, 1.0].min
+        end
+      end
+
+      def migration_class
+        BackgroundMigration.named(migration_name)
+      end
+
+      def migration_object
+        @migration_object ||= migration_class.new(*arguments)
+      end
+
+      def migration_relation
+        migration_object.relation
+      end
+
+      # Returns whether the interval between previous step run has passed.
+      # @return [Boolean]
+      #
+      def interval_elapsed?
+        if migration_jobs.running.exists?
+          false
+        elsif (job = last_completed_job)
+          job.finished_at + batch_pause <= Time.current
+        else
+          true
+        end
+      end
+
+      # Manually retry failed jobs.
+      #
+      # This method marks failed jobs as ready to be processed again, and
+      # they will be picked up on the next Scheduler run.
+      #
+      def retry_failed_jobs
+        iterator = BatchIterator.new(migration_jobs.failed)
+        iterator.each_batch(of: 100) do |batch|
+          transaction do
+            batch.each(&:retry)
+            enqueued!
+          end
+        end
+      end
+
+      # @private
+      def next_batch_range
+        iterator = BatchIterator.new(migration_relation)
+        batch_range = nil
+
+        # rubocop:disable Lint/UnreachableLoop
+        iterator.each_batch(of: batch_size, column: batch_column_name, start: next_min_value) do |relation|
+          if Utils.ar_version <= 4.2
+            # ActiveRecord <= 4.2 does not support pluck with Arel nodes
+            quoted_column = self.class.connection.quote_column_name(batch_column_name)
+            batch_range = relation.pluck("MIN(#{quoted_column}), MAX(#{quoted_column})").first
+          else
+            min = relation.arel_table[batch_column_name].minimum
+            max = relation.arel_table[batch_column_name].maximum
+
+            batch_range = relation.pluck(min, max).first
+          end
+          break
+        end
+        # rubocop:enable Lint/UnreachableLoop
+
+        return if batch_range.nil?
+
+        min_value, max_value = batch_range
+        return if min_value > self.max_value
+
+        max_value = [max_value, self.max_value].min
+
+        [min_value, max_value]
+      end
+
+      private
+        def validate_batch_column_values
+          if max_value.to_i < min_value.to_i
+            errors.add(:base, "max_value should be greater than or equal to min_value")
+          end
+        end
+
+        def validate_batch_sizes
+          if sub_batch_size.to_i > batch_size.to_i
+            errors.add(:base, "sub_batch_size should be smaller than or equal to batch_size")
+          end
+        end
+
+        def validate_jobs_status
+          if succeeded? && migration_jobs.except_succeeded.exists?
+            errors.add(:base, "all migration jobs must be succeeded")
+          elsif failed? && !migration_jobs.failed.exists?
+            errors.add(:base, "at least one migration job must be failed")
+          end
+        end
+
+        def set_defaults
+          if migration_relation.is_a?(ActiveRecord::Relation)
+            self.batch_column_name  ||= migration_relation.primary_key
+            self.min_value          ||= migration_relation.minimum(batch_column_name)
+            self.max_value          ||= migration_relation.maximum(batch_column_name)
+
+            count = migration_object.count
+            self.rows_count = count if count != :no_count
+          end
+
+          config = ::OnlineMigrations.config.background_migrations
+          self.batch_size           ||= config.batch_size
+          self.sub_batch_size       ||= config.sub_batch_size
+          self.batch_pause          ||= config.batch_pause
+          self.sub_batch_pause_ms   ||= config.sub_batch_pause_ms
+          self.batch_max_attempts   ||= config.batch_max_attempts
+
+          # This can be the case when run in development on empty tables
+          if min_value.nil?
+            # integer IDs minimum value is 1
+            self.min_value = self.max_value = 1
+          end
+        end
+
+        def next_min_value
+          if last_job
+            last_job.max_value.next
+          else
+            min_value
+          end
+        end
+    end
+  end
+end

data/lib/online_migrations/background_migrations/migration_helpers.rb

@@ -0,0 +1,238 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundMigrations
+    module MigrationHelpers
+      # Backfills column data using background migrations.
+      #
+      # @param table_name [String, Symbol]
+      # @param column_name [String, Symbol]
+      # @param value
+      # @param model_name [String] If Active Record multiple databases feature is used,
+      #     the class name of the model to get connection from.
+      # @param options [Hash] used to control the behavior of background migration.
+      #     See `#enqueue_background_migration`
+      #
+      # @return [OnlineMigrations::BackgroundMigrations::Migration]
+      #
+      # @example
+      #   backfill_column_in_background(:users, :admin, false)
+      #
+      # @example Additional background migration options
+      #   backfill_column_in_background(:users, :admin, false, batch_size: 10_000)
+      #
+      # @note This method is better suited for extra large tables (100s of millions of records).
+      #     For smaller tables it is probably better and easier to use more flexible `update_column_in_batches`.
+      #
+      # @note Consider `backfill_columns_in_background` when backfilling multiple columns
+      #   to avoid rewriting the table multiple times.
+      #
+      def backfill_column_in_background(table_name, column_name, value, model_name: nil, **options)
+        backfill_columns_in_background(table_name, { column_name => value },
+                                       model_name: model_name, **options)
+      end
+
+      # Same as `backfill_column_in_background` but for multiple columns.
+      #
+      # @param updates [Hash] keys - column names, values - corresponding values
+      #
+      # @example
+      #   backfill_columns_in_background(:users, { admin: false, status: "active" })
+      #
+      # @see #backfill_column_in_background
+      #
+      def backfill_columns_in_background(table_name, updates, model_name: nil, **options)
+        model_name = model_name.name if model_name.is_a?(Class)
+
+        enqueue_background_migration(
+          "BackfillColumn",
+          table_name,
+          updates,
+          model_name,
+          **options
+        )
+      end
+
+      # Backfills data from the old column to the new column using background migrations.
+      #
+      # @param table_name [String, Symbol]
+      # @param column_name [String, Symbol]
+      # @param model_name [String] If Active Record multiple databases feature is used,
+      #     the class name of the model to get connection from.
+      # @param type_cast_function [String, Symbol] Some type changes require casting data to a new type.
+      #     For example when changing from `text` to `jsonb`. In this case, use the `type_cast_function` option.
+      #     You need to make sure there is no bad data and the cast will always succeed
+      # @param options [Hash] used to control the behavior of background migration.
+      #     See `#enqueue_background_migration`
+      #
+      # @return [OnlineMigrations::BackgroundMigrations::Migration]
+      #
+      # @example
+      #   backfill_column_for_type_change_in_background(:files, :size)
+      #
+      # @example With type casting
+      #   backfill_column_for_type_change_in_background(:users, :settings, type_cast_function: "jsonb")
+      #
+      # @example Additional background migration options
+      #   backfill_column_for_type_change_in_background(:files, :size, batch_size: 10_000)
+      #
+      # @note This method is better suited for extra large tables (100s of millions of records).
+      #     For smaller tables it is probably better and easier to use more flexible `backfill_column_for_type_change`.
+      #
+      def backfill_column_for_type_change_in_background(table_name, column_name, model_name: nil,
+                                                        type_cast_function: nil, **options)
+        backfill_columns_for_type_change_in_background(
+          table_name,
+          column_name,
+          model_name: model_name,
+          type_cast_functions: { column_name => type_cast_function },
+          **options
+        )
+      end
+
+      # Same as `backfill_column_for_type_change_in_background` but for multiple columns.
+      #
+      # @param type_cast_functions [Hash] if not empty, keys - column names,
+      #   values - corresponding type cast functions
+      #
+      # @see #backfill_column_for_type_change_in_background
+      #
+      def backfill_columns_for_type_change_in_background(table_name, *column_names, model_name: nil,
+                                                         type_cast_functions: {}, **options)
+        tmp_columns = column_names.map { |column_name| "#{column_name}_for_type_change" }
+        model_name = model_name.name if model_name.is_a?(Class)
+
+        enqueue_background_migration(
+          "CopyColumn",
+          table_name,
+          column_names,
+          tmp_columns,
+          model_name,
+          type_cast_functions,
+          **options
+        )
+      end
+
+      # Copies data from the old column to the new column using background migrations.
+      #
+      # @param table_name [String, Symbol]
+      # @param copy_from [String, Symbol] source column name
+      # @param copy_to [String, Symbol] destination column name
+      # @param model_name [String] If Active Record multiple databases feature is used,
+      #     the class name of the model to get connection from.
+      # @param type_cast_function [String, Symbol] Some type changes require casting data to a new type.
+      #     For example when changing from `text` to `jsonb`. In this case, use the `type_cast_function` option.
+      #     You need to make sure there is no bad data and the cast will always succeed
+      # @param options [Hash] used to control the behavior of background migration.
+      #     See `#enqueue_background_migration`
+      #
+      # @return [OnlineMigrations::BackgroundMigrations::Migration]
+      #
+      # @example
+      #   copy_column_in_background(:users, :id, :id_for_type_change)
+      #
+      # @note This method is better suited for extra large tables (100s of millions of records).
+      #     For smaller tables it is probably better and easier to use more flexible `update_column_in_batches`.
+      #
+      def copy_column_in_background(table_name, copy_from, copy_to, model_name: nil, type_cast_function: nil, **options)
+        copy_columns_in_background(
+          table_name,
+          [copy_from],
+          [copy_to],
+          model_name: model_name,
+          type_cast_functions: { copy_from => type_cast_function },
+          **options
+        )
+      end
+
+      # Same as `copy_column_in_background` but for multiple columns.
+      #
+      # @param type_cast_functions [Hash] if not empty, keys - column names,
+      #   values - corresponding type cast functions
+      #
+      # @see #copy_column_in_background
+      #
+      def copy_columns_in_background(table_name, copy_from, copy_to, model_name: nil, type_cast_functions: {}, **options)
+        model_name = model_name.name if model_name.is_a?(Class)
+
+        enqueue_background_migration(
+          "CopyColumn",
+          table_name,
+          copy_from,
+          copy_to,
+          model_name,
+          type_cast_functions,
+          **options
+        )
+      end
+
+      # Creates a background migration for the given job class name.
+      #
+      # A background migration runs one job at a time, computing the bounds of the next batch
+      # 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 (20_000) Number of rows to process in a single background migration run
+      # @option options [Integer] :sub_batch_size (1000) 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]
+      #
+      # @example
+      #   enqueue_background_migration("BackfillProjectIssuesCount",
+      #       batch_size: 10_000, batch_max_attempts: 10)
+      #
+      #   # Given the background migration exists:
+      #
+      #   class BackfillProjectIssuesCount < OnlineMigrations::BackgroundMigration
+      #     def relation
+      #       Project.all
+      #     end
+      #
+      #     def process_batch(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
+      #     def count
+      #       Project.maximum(:id)
+      #     end
+      #   end
+      #
+      # @note For convenience, the enqueued background migration is run inline
+      #     in development and test environments
+      #
+      def enqueue_background_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)
+
+        migration_name = migration_name.name if migration_name.is_a?(Class)
+
+        migration = Migration.create!(
+          migration_name: migration_name,
+          arguments: arguments,
+          **options
+        )
+
+        # For convenience in dev/test environments
+        if Utils.developer_env?
+          runner = MigrationRunner.new(migration)
+          runner.run_all_migration_jobs
+        end
+
+        migration
+      end
+    end
+  end
+end

data/lib/online_migrations/background_migrations/migration_job.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundMigrations
+    class MigrationJob < ActiveRecord::Base
+      STATUSES = [
+        :enqueued,
+        :running,
+        :failed,
+        :succeeded,
+      ]
+
+      self.table_name = :background_migration_jobs
+
+      # For ActiveRecord <= 4.2 needs to fully specify enum values
+      scope :active, -> { where(status: [statuses[:enqueued], statuses[:running]]) }
+      scope :completed, -> { where(status: [statuses[:failed], statuses[:succeeded]]) }
+      scope :stuck, -> do
+        timeout = ::OnlineMigrations.config.background_migrations.stuck_jobs_timeout
+        active.where("updated_at <= ?", timeout.ago)
+      end
+
+      scope :retriable, -> do
+        failed_retriable = 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
+
+      scope :except_succeeded, -> { where("status != ?", statuses[:succeeded]) }
+
+      enum status: STATUSES.map { |status| [status, status.to_s] }.to_h
+
+      delegate :migration_class, :migration_object, :migration_relation, :batch_column_name,
+        :arguments, :batch_pause, to: :migration
+
+      belongs_to :migration
+
+      # For ActiveRecord 5.0+ this is validated by default from belongs_to
+      validates :migration, presence: true
+
+      validates :min_value, :max_value, presence: true, numericality: { greater_than: 0 }
+      validate :values_in_migration_range, if: :min_value?
+      validate :validate_values_order, if: :min_value?
+
+      validates_with MigrationJobStatusValidator, on: :update
+
+      before_create :copy_settings_from_migration
+
+      # Mark this job as ready to be processed again.
+      #
+      # This is used when retrying failed jobs.
+      #
+      def retry
+        update!(
+          status: self.class.statuses[:enqueued],
+          attempts: 0,
+          started_at: nil,
+          finished_at: nil
+        )
+      end
+
+      private
+        def values_in_migration_range
+          if min_value < migration.min_value || max_value > migration.max_value
+            errors.add(:base, "min_value and max_value should be in background migration values range")
+          end
+        end
+
+        def validate_values_order
+          if max_value.to_i < min_value.to_i
+            errors.add(:base, "max_value should be greater than or equal to min_value")
+          end
+        end
+
+        def copy_settings_from_migration
+          self.batch_size       = migration.batch_size
+          self.sub_batch_size   = migration.sub_batch_size
+          self.pause_ms         = migration.sub_batch_pause_ms
+          self.max_attempts     = migration.batch_max_attempts
+        end
+    end
+  end
+end

data/lib/online_migrations/background_migrations/migration_job_runner.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+module OnlineMigrations
+  module BackgroundMigrations
+    # @private
+    class MigrationJobRunner
+      attr_reader :migration_job
+
+      delegate :attempts, :migration_relation, :migration_object, :sub_batch_size,
+        :batch_column_name, :min_value, :max_value, :pause_ms, to: :migration_job
+
+      def initialize(migration_job)
+        @migration_job = migration_job
+      end
+
+      def run
+        job_payload = { background_migration_job: migration_job }
+        if migration_job.attempts >= 1
+          ActiveSupport::Notifications.instrument("retried.background_migrations", job_payload)
+        end
+
+        migration_job.update!(
+          attempts: attempts + 1,
+          status: :running,
+          started_at: Time.current,
+          finished_at: nil,
+          error_class: nil,
+          error_message: nil,
+          backtrace: nil
+        )
+
+        ActiveSupport::Notifications.instrument("process_batch.background_migrations", job_payload) do
+          run_batch
+        end
+
+        migration_job.update!(status: :succeeded, finished_at: Time.current)
+      rescue Exception => e # rubocop:disable Lint/RescueException
+        backtrace_cleaner = ::OnlineMigrations.config.background_migrations.backtrace_cleaner
+
+        migration_job.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_migrations.error_handler.call(e, migration_job)
+      end
+
+      private
+        def run_batch
+          iterator = ::OnlineMigrations::BatchIterator.new(migration_relation)
+
+          iterator.each_batch(of: sub_batch_size, column: batch_column_name,
+                              start: min_value, finish: max_value) do |sub_batch|
+            migration_object.process_batch(sub_batch)
+            sleep(pause_ms * 0.001) if pause_ms > 0
+          end
+        end
+    end
+  end
+end