data_shifter 0.1.0

data/lib/data_shifter/rubocop.rb

@@ -0,0 +1,4 @@
+# frozen_string_literal: true
+
+require "rubocop"
+require "rubocop/cop/data_shifter/skip_transaction_guard_dry_run"

data/lib/data_shifter/shift.rb

@@ -0,0 +1,373 @@
+# frozen_string_literal: true
+
+require "axn"
+require "active_support/isolated_execution_state"
+require_relative "internal/env"
+require_relative "internal/output"
+require_relative "internal/signal_handler"
+require_relative "internal/record_utils"
+require_relative "internal/progress_bar"
+
+# Base class for data shifts. Dry-run by default, progress bars, transaction modes, consistent summaries.
+#
+# Usage:
+#
+#   # lib/data_shifts/20260201120000_backfill_foo.rb
+#   module DataShifts
+#     class BackfillFoo < DataShifter::Shift
+#       description "Backfill foo on bars"
+#
+#       def collection
+#         Bar.where(foo: nil)
+#       end
+#
+#       def process_record(bar)
+#         bar.update!(foo: computed_value(bar))
+#       end
+#     end
+#   end
+#
+# Running:
+#   - `rake data:shift:backfill_foo` (dry run by default)
+#   - `COMMIT=1 rake data:shift:backfill_foo` (apply changes)
+#   - Or call directly: `MyShift.call(dry_run: false)` (Axn semantics) - but note default location not auto-loaded
+#
+# Transaction modes (set at class level with `transaction`):
+#   - `transaction :single` (default): one transaction for the whole run (all-or-nothing).
+#   - `transaction :per_record`: each record in its own transaction.
+#   - `transaction false`: no automatic transactions; guard writes with `return if dry_run?`.
+#
+# Dry run: In `:single` and `:per_record`, dry_run rolls back DB changes automatically.
+# Non-DB side effects are not rolled back; guard with `return if dry_run?` / `return unless dry_run?`.
+#
+# Fixed list of IDs (fail fast): Use find_exactly!(Model, [id1, id2, ...]) in `collection`.
+# Large collections: Return an ActiveRecord::Relation and iteration uses `find_each`.
+#
+module DataShifter
+  class Shift
+    include Axn
+
+    expects :dry_run, type: :boolean, default: true
+
+    log_calls false if respond_to?(:log_calls)
+
+    around :_with_transaction_for_dry_run
+    before :_reset_tracking
+    on_success :_print_summary
+    on_error :_print_summary
+
+    class_attribute :_transaction_mode, default: :single
+    class_attribute :_progress_enabled, default: true
+    class_attribute :_description, default: nil
+    class_attribute :_task_name, default: nil
+    class_attribute :_throttle_interval, default: nil
+
+    class << self
+      def description(text = nil)
+        if text.nil?
+          _description
+        else
+          self._description = text.to_s.presence
+        end
+      end
+
+      def task_name(value = nil)
+        if value.nil?
+          _task_name
+        else
+          self._task_name = value.to_s.presence
+        end
+      end
+
+      def transaction(mode)
+        case mode
+        when :per_record
+          self._transaction_mode = :per_record
+        when :none, false
+          self._transaction_mode = :none
+        when :single, true
+          self._transaction_mode = :single
+        else
+          raise ArgumentError, "Invalid transaction mode: #{mode.inspect}. Expected :single, :per_record, :none, true, or false."
+        end
+      end
+
+      def progress(enabled = nil)
+        if enabled.nil?
+          _progress_enabled
+        else
+          self._progress_enabled = !!enabled
+        end
+      end
+
+      def throttle(interval)
+        self._throttle_interval = interval
+      end
+
+      def run!
+        dry_run = Internal::Env.dry_run?
+        result = call(dry_run:)
+        raise result.exception if result.exception
+        raise StandardError, result.error unless result.ok?
+      end
+    end
+
+    # --- Public API (intentionally exposed to subclasses) ---
+
+    def call
+      _for_each_record_in(collection) { |record| process_record(record) }
+    end
+
+    def find_exactly!(model, ids)
+      ids = Array(ids).compact.uniq
+      return model.none if ids.empty?
+
+      records_by_id = model.where(id: ids).index_by(&:id)
+      missing = ids.reject { |id| records_by_id.key?(id) }
+      raise "Expected #{model.name} with ids #{ids.inspect}, but missing: #{missing.inspect}" if missing.any?
+
+      ids.map { |id| records_by_id[id] }
+    end
+
+    def dry_run? = dry_run
+
+    def skip!(reason = nil)
+      @stats[:skipped] += 1
+      @stats[:succeeded] -= 1
+      log "  SKIP: #{reason}" if reason
+    end
+
+    def log(message)
+      puts message
+    end
+
+    private
+
+    # --- Axn lifecycle hooks ---
+
+    def _with_transaction_for_dry_run(chain)
+      if _transaction_mode == :none
+        chain.call
+        return
+      end
+
+      if _transaction_mode == :single
+        ActiveRecord::Base.transaction do
+          chain.call
+          raise ActiveRecord::Rollback if dry_run?
+        end
+        return
+      end
+
+      if dry_run?
+        ActiveRecord::Base.transaction do
+          chain.call
+          raise ActiveRecord::Rollback
+        end
+      else
+        chain.call
+      end
+    end
+
+    def _reset_tracking
+      @stats = { processed: 0, succeeded: 0, failed: 0, skipped: 0 }
+      @errors = []
+      @start_time = Time.current
+      @last_status_print = @start_time
+      @_data_shift_interrupted = false
+      @_last_successful_id = nil
+    end
+
+    def _print_summary
+      Internal::Output.print_summary(
+        io: $stdout,
+        stats: @stats,
+        errors: @errors,
+        start_time: @start_time,
+        dry_run: dry_run?,
+        transaction_mode: _transaction_mode,
+        interrupted: @_data_shift_interrupted,
+        task_name: self.class.task_name,
+        last_successful_id: @_last_successful_id,
+      )
+    end
+
+    # --- Override points ---
+
+    def collection
+      raise NotImplementedError, "#{self.class.name}: override `collection`"
+    end
+
+    def process_record(_record)
+      raise NotImplementedError, "#{self.class.name}: override `process_record`"
+    end
+
+    # --- Record iteration ---
+
+    def _print_progress
+      Internal::Output.print_progress(
+        io: $stdout,
+        stats: @stats,
+        errors: @errors,
+        start_time: @start_time,
+        status_interval: Internal::Env.status_interval_seconds,
+      )
+    end
+
+    def _for_each_record_in(records, label: nil, &)
+      _reset_tracking
+      ActiveSupport::IsolatedExecutionState[:_data_shifter_current_run] = self
+      status_proc = proc { ActiveSupport::IsolatedExecutionState[:_data_shifter_current_run]&.send(:_print_progress) }
+      prev_handlers = Internal::SignalHandler.install_status_traps(status_proc)
+      begin
+        _each_record_impl(records, label:, &)
+      rescue Interrupt
+        _handle_interrupt
+      ensure
+        ActiveSupport::IsolatedExecutionState.delete(:_data_shifter_current_run)
+        Internal::SignalHandler.restore_status_traps(prev_handlers)
+      end
+    end
+
+    def _each_record_impl(records, label: nil, &)
+      records = _apply_continue_from(records)
+
+      if records.respond_to?(:find_each)
+        total = records.count
+        @label = label || Internal::RecordUtils.default_label_for_relation(records)
+        _print_header(total)
+        enum = records
+      else
+        items = records.respond_to?(:to_a) ? records.to_a : Array(records)
+        total = items.size
+        @label = label || Internal::RecordUtils.default_label(items)
+        _print_header(total)
+        enum = items
+      end
+
+      case _transaction_mode
+      when :single
+        _run_in_single_transaction(enum, total, &)
+      when :per_record
+        _run_per_record(enum, total, &)
+      when :none
+        _run_without_transaction(enum, total, &)
+      end
+
+      fail! "#{@stats[:failed]} record(s) failed" if @errors.any?
+    end
+
+    def _apply_continue_from(records)
+      continue_from = Internal::Env.continue_from_id
+      return records if continue_from.nil?
+
+      unless records.respond_to?(:find_each)
+        raise ArgumentError,
+              "CONTINUE_FROM is only supported for ActiveRecord::Relation collections. " \
+              "Array-based collections (e.g. from find_exactly!) cannot be resumed."
+      end
+
+      primary_key = records.model.primary_key
+      log "[CONTINUE_FROM] Resuming from #{primary_key} > #{continue_from}"
+      records.where("#{records.model.quoted_table_name}.#{primary_key} > ?", continue_from)
+    end
+
+    # --- Transaction execution strategies ---
+
+    def _run_in_single_transaction(enum, total, &block)
+      ActiveRecord::Base.transaction do
+        _iterate(enum, total, &block)
+        if dry_run?
+          log "\nDry run complete — rolling back all changes."
+          raise ActiveRecord::Rollback
+        end
+      end
+    rescue StandardError => e
+      return if @errors.any?
+
+      @stats[:failed] += 1
+      @errors << { record: "transaction", error: e.message, backtrace: e.backtrace&.first(3) }
+    end
+
+    def _run_per_record(enum, total, &)
+      _iterate(enum, total) do |record|
+        if dry_run?
+          yield record
+        else
+          ActiveRecord::Base.transaction { yield record }
+        end
+      end
+    end
+
+    def _run_without_transaction(enum, total, &)
+      _iterate(enum, total, &)
+    end
+
+    def _iterate(enum, total)
+      bar = Internal::ProgressBar.create(total:, dry_run: dry_run?, enabled: _progress_enabled)
+      if enum.respond_to?(:find_each)
+        enum.find_each do |record|
+          _process_one(record) { yield record }
+          bar&.increment
+          sleep(_throttle_interval) if _throttle_interval
+        end
+      else
+        enum.each do |record|
+          _process_one(record) { yield record }
+          bar&.increment
+          sleep(_throttle_interval) if _throttle_interval
+        end
+      end
+    end
+
+    def _process_one(record)
+      @stats[:processed] += 1
+      yield
+      @stats[:succeeded] += 1
+      @_last_successful_id = record.id if record.respond_to?(:id)
+    rescue StandardError => e
+      @stats[:failed] += 1
+      identifier = Internal::RecordUtils.identifier(record)
+      @errors << { record: identifier, error: e.message, backtrace: e.backtrace&.first(3) }
+      log "ERROR #{identifier}: #{e.message}"
+
+      raise if _transaction_mode == :single
+    ensure
+      _maybe_print_interval_status
+    end
+
+    def _maybe_print_interval_status
+      interval = Internal::Env.status_interval_seconds
+      return unless interval&.positive?
+      return unless @start_time && (Time.current - @last_status_print) >= interval
+
+      @last_status_print = Time.current
+      _print_progress
+    end
+
+    # --- Output helpers ---
+
+    def _print_header(total)
+      Internal::Output.print_header(
+        io: $stdout,
+        shift_class: self.class,
+        total:,
+        label: @label,
+        dry_run: dry_run?,
+        transaction_mode: _transaction_mode,
+        status_interval: Internal::Env.status_interval_seconds,
+      )
+    end
+
+    def _handle_interrupt
+      @_data_shift_interrupted = true
+      log "\n\n*** Interrupted by user (Ctrl+C) ***"
+
+      # Print summary now since on_error may not fire for Interrupt (SignalException)
+      _print_summary
+
+      # Re-raise to trigger transaction rollback in the wrapping transaction block
+      raise Interrupt
+    end
+  end
+end

data/lib/data_shifter/spec_helper.rb

@@ -0,0 +1,75 @@
+# frozen_string_literal: true
+
+module DataShifter
+  # Test helpers for RSpec. Include this module in your spec_helper or rails_helper:
+  #
+  #   require "data_shifter/spec_helper"
+  #
+  #   RSpec.configure do |config|
+  #     config.include DataShifter::SpecHelper, type: :data_shift
+  #   end
+  #
+  # Or include it in individual specs:
+  #
+  #   RSpec.describe DataShifts::BackfillFoo do
+  #     include DataShifter::SpecHelper
+  #     ...
+  #   end
+  #
+  module SpecHelper
+    # Run a data shift class with the given options.
+    # Returns the Axn::Result.
+    #
+    # @param shift_class [Class] the DataShifter::Shift subclass
+    # @param dry_run [Boolean] whether to run in dry_run mode (default: true)
+    # @param commit [Boolean] shorthand for dry_run: false (default: false)
+    # @return [Axn::Result]
+    #
+    # @example
+    #   result = run_data_shift(DataShifts::BackfillFoo)
+    #   expect(result).to be_ok
+    #
+    # @example with commit
+    #   result = run_data_shift(DataShifts::BackfillFoo, commit: true)
+    #   expect(record.reload.foo).to eq("bar")
+    #
+    def run_data_shift(shift_class, dry_run: true, commit: false)
+      effective_dry_run = commit ? false : dry_run
+      shift_class.call(dry_run: effective_dry_run)
+    end
+
+    # Suppress STDOUT output during a block (useful for cleaner test output).
+    #
+    # @example
+    #   silence_data_shift_output do
+    #     run_data_shift(DataShifts::BackfillFoo, commit: true)
+    #   end
+    #
+    def silence_data_shift_output
+      original_stdout = $stdout
+      $stdout = StringIO.new
+      yield
+    ensure
+      $stdout = original_stdout
+    end
+
+    # Run a shift and capture its output.
+    # Returns [Axn::Result, String] tuple.
+    #
+    # @example
+    #   result, output = capture_data_shift_output do
+    #     run_data_shift(DataShifts::BackfillFoo)
+    #   end
+    #   expect(output).to include("DRY RUN")
+    #
+    def capture_data_shift_output
+      original_stdout = $stdout
+      $stdout = StringIO.new
+      result = yield
+      output = $stdout.string
+      [result, output]
+    ensure
+      $stdout = original_stdout
+    end
+  end
+end

data/lib/data_shifter/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module DataShifter
+  VERSION = "0.1.0"
+end

data/lib/data_shifter.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+require_relative "data_shifter/version"
+require_relative "data_shifter/shift"
+require_relative "data_shifter/railtie"

data/lib/generators/data_shift_generator.rb

@@ -0,0 +1,132 @@
+# frozen_string_literal: true
+
+# Generator for data shifts.
+#
+# Usage:
+#   rails g data_shift backfill_users
+#   rails g data_shift backfill_users --model=User
+#
+class DataShiftGenerator < Rails::Generators::NamedBase
+  class_option :model,
+               type: :string,
+               default: nil,
+               desc: "Model to operate on (e.g. User). Pre-fills the collection method."
+
+  class_option :spec,
+               type: :boolean,
+               default: false,
+               desc: "Generate RSpec file"
+
+  def check_for_naming_conflict
+    underscored_name = name.underscore
+
+    # Use destination_root if available (for testing), otherwise Rails.root
+    root = respond_to?(:destination_root) ? Pathname.new(destination_root) : Rails.root
+    shifts_dir = root.join("lib/data_shifts")
+    return unless shifts_dir.exist?
+
+    # Look for any existing file that would create the same rake task name
+    conflicting_file = Dir.glob(shifts_dir.join("*_#{underscored_name}.rb")).first
+    return unless conflicting_file
+
+    raise Thor::Error, <<~ERROR
+      A data shift with task name '#{underscored_name}' already exists:
+        #{conflicting_file}
+
+      Rake task names must be unique. Please choose a different name.
+    ERROR
+  end
+
+  def create_shift_file
+    underscored_name = name.underscore
+    @timestamp = Time.current.strftime("%Y%m%d%H%M%S")
+    @class_name = underscored_name.camelize
+    model_name_raw = options[:model].to_s.strip
+    @model_name = model_name_raw.present? ? model_name_raw.underscore.singularize.camelize : nil
+
+    collection_body = if @model_name.present?
+                        "#{@model_name}.all"
+                      else
+                        "# Model.where(foo: nil)"
+                      end
+
+    record_arg = @model_name.present? ? @model_name.underscore : "record"
+
+    create_file "lib/data_shifts/#{@timestamp}_#{underscored_name}.rb", <<~RUBY
+      # frozen_string_literal: true
+
+      #   rake data:shift:#{underscored_name}          # Dry run (default)
+      #   COMMIT=1 rake data:shift:#{underscored_name} # Apply changes
+
+      module DataShifts
+        class #{@class_name} < DataShifter::Shift
+          description "TODO: Describe this shift"
+
+          transaction true # or false or :per_record
+
+          def collection
+            #{collection_body}
+          end
+
+          def process_record(#{record_arg})
+            # #{record_arg}.update!(...)
+          end
+        end
+      end
+    RUBY
+  end
+
+  def create_spec_file
+    return unless options[:spec]
+    return unless rspec_enabled?
+
+    underscored_name = name.underscore
+    record_arg = @model_name.present? ? @model_name.underscore : "record"
+
+    create_file "spec/lib/data_shifts/#{underscored_name}_spec.rb", <<~RUBY
+      # frozen_string_literal: true
+
+      require "rails_helper"
+      require "data_shifter/spec_helper"
+
+      RSpec.describe DataShifts::#{@class_name} do
+        include DataShifter::SpecHelper
+
+        before { allow($stdout).to receive(:puts) }
+
+        # TODO: Set up test records
+        # let(:#{record_arg}) { create(:#{record_arg}) }
+
+        describe "dry run" do
+          it "does not persist changes" do
+            result = run_data_shift(described_class, dry_run: true)
+            expect(result).to be_ok
+            # TODO: Assert that records are unchanged
+          end
+        end
+
+        describe "commit" do
+          it "applies changes" do
+            result = run_data_shift(described_class, commit: true)
+            expect(result).to be_ok
+            # TODO: Assert that records are updated
+          end
+        end
+      end
+    RUBY
+  end
+
+  private
+
+  def rspec_enabled?
+    # Check if rspec-rails is available and configured as the test framework
+    return false unless defined?(Rails)
+
+    # Check Rails generator configuration
+    test_framework = Rails.configuration.generators.options.dig(:rails, :test_framework)
+    return test_framework == :rspec if test_framework
+
+    # Fall back to checking if rspec-rails is loaded
+    defined?(RSpec::Rails)
+  end
+end

data/lib/rubocop/cop/data_shifter/skip_transaction_guard_dry_run.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+module RuboCop
+  module Cop
+    module DataShifter
+      # In data shift files, `transaction false` disables automatic transaction
+      # and rollback. DB writes (and side effects) are not rolled back on dry run, so
+      # the shift must guard them with `return if dry_run?` or `return unless dry_run?`.
+      #
+      # @example
+      #   # bad
+      #   class BackfillUsers < DataShifter::Shift
+      #     transaction false
+      #     def process_record(record)
+      #       record.update!(foo: 1)
+      #     end
+      #   end
+      #
+      #   # good
+      #   class BackfillUsers < DataShifter::Shift
+      #     transaction false
+      #     def process_record(record)
+      #       return if dry_run?
+      #       record.update!(foo: 1)
+      #     end
+      #   end
+      class SkipTransactionGuardDryRun < Base
+        MSG = "Data shifts using `transaction false` must guard writes/side effects with " \
+              "`return if dry_run?` or `return unless dry_run?`."
+
+        def_node_matcher :skip_transaction_call?, <<~PATTERN
+          (send _ :transaction {(sym :none) (false)})
+        PATTERN
+
+        def on_send(node)
+          return unless skip_transaction_call?(node)
+          return if file_contains_dry_run_guard?
+
+          add_offense(node, message: MSG)
+        end
+
+        private
+
+        def file_contains_dry_run_guard?
+          return true unless processed_source.ast
+
+          processed_source.ast.each_node(:send) do |send_node|
+            return true if send_node.method?(:dry_run?)
+          end
+          false
+        end
+      end
+    end
+  end
+end