cold_storage 0.1.0

data/lib/cold_storage/archive_model.rb

@@ -0,0 +1,181 @@
+# frozen_string_literal: true
+
+require 'monitor'
+
+module ColdStorage
+  # Namespace holding the generated archive classes, so they have real names in
+  # logs and backtraces (ColdStorage::Archived::Invoices).
+  #
+  # Classes appear here on demand: an association declared on one archive model
+  # names its counterpart, and the constant is built the first time it is
+  # touched.
+  module Archived
+    def self.const_missing(name)
+      source = ArchiveModel.registered_source(name)
+      return super if source.nil?
+
+      ArchiveModel.for(source)
+    end
+  end
+
+  # Builds, and caches, the ActiveRecord class that reads and writes a mirrored
+  # table in the archive database.
+  module ArchiveModel
+    MIRRORED_MACROS = %i[has_many has_one belongs_to].freeze
+    MONITOR = Monitor.new
+    private_constant :MONITOR
+
+    class << self
+      # @param source [Class] an ActiveRecord model in the primary database
+      # @return [Class] its counterpart in the archive database
+      def for(source)
+        ArchiveRecord.connect!
+        cache[source.table_name] || MONITOR.synchronize { cache[source.table_name] ||= build(source) }
+      end
+
+      # @api private the source model behind a generated constant name
+      def registered_source(const_name)
+        MONITOR.synchronize { sources[const_name.to_s] }
+      end
+
+      def clear!
+        MONITOR.synchronize do
+          @cache = {}
+          @sources = {}
+        end
+      end
+
+      private
+
+      def cache
+        @cache ||= {}
+      end
+
+      # Constant name => source model, so Archived.const_missing can build a
+      # class an association asked for.
+      def sources
+        @sources ||= {}
+      end
+
+      def build(source)
+        table = source.table_name
+        polymorphic = source.polymorphic_name
+
+        klass = Class.new(ArchiveRecord) do
+          self.table_name         = table
+          self.inheritance_column = nil
+          self.record_timestamps  = false
+        end
+        klass.primary_key = source.primary_key
+        klass.source_model_name = source.name
+        # So `has_many ..., as:` conditions keep matching the archived rows,
+        # which store the *source* class name in their type column.
+        klass.define_singleton_method(:polymorphic_name) { polymorphic }
+
+        register(source)
+        const = const_name_for(table)
+        Archived.send(:remove_const, const) if Archived.const_defined?(const, false)
+        Archived.const_set(const, klass)
+        mirror_associations(klass, source)
+        klass
+      end
+
+      def register(source)
+        sources[const_name_for(source.table_name)] = source
+      end
+
+      # Redeclares the source associations against the archive classes, so an
+      # archived row can be read together with its archived children.
+      #
+      # Deliberately dropped: dependent:, counter caches, touch and autosave -
+      # an archive is a reading surface, it must not cascade anything.
+      def mirror_associations(klass, source)
+        source.reflect_on_all_associations.each do |reflection|
+          next unless MIRRORED_MACROS.include?(reflection.macro)
+          next if reflection.options[:polymorphic]
+
+          if reflection.options[:through]
+            mirror_through_association(klass, reflection)
+          else
+            mirror_association(klass, reflection)
+          end
+        end
+      end
+
+      # A through association needs no class_name: it composes out of the two
+      # associations it is built on, which are mirrored too, so it stays inside
+      # the archive database.
+      def mirror_through_association(klass, reflection)
+        # A polymorphic source would resolve to a primary-database class and
+        # quietly read from the wrong database.
+        return if reflection.options[:source_type]
+
+        options = { through: reflection.options[:through] }
+        options[:source] = reflection.options[:source] if reflection.options[:source]
+
+        klass.public_send(reflection.macro, reflection.name, reflection.scope, **options)
+      rescue StandardError => e
+        log_skipped(klass, reflection, e)
+      end
+
+      def mirror_association(klass, reflection)
+        target = reflection.klass
+        return if target.abstract_class? || target.table_name.blank?
+
+        register(target)
+        klass.public_send(reflection.macro, reflection.name, portable_scope(reflection),
+                          **options_for(reflection, target))
+      rescue StandardError => e
+        log_skipped(klass, reflection, e)
+      end
+
+      # An association scope is written against the source model, and may call
+      # scopes or class methods only that model has:
+      #
+      #   has_many :time_periods, -> { sorted('start_time', 'asc') }
+      #
+      # Plain conditions carry over to the archive; anything the archive class
+      # cannot evaluate falls back to the unfiltered relation rather than
+      # blowing up a read.
+      def portable_scope(reflection)
+        scope = reflection.scope
+        return nil if scope.nil?
+
+        lambda do |*args|
+          instance_exec(*args, &scope)
+        rescue StandardError => e
+          ColdStorage.logger.debug do
+            "#{Logging::PREFIX} ignoring the scope of #{reflection.name} on the archive model: " \
+              "#{e.class}: #{e.message}"
+          end
+          self
+        end
+      end
+
+      def log_skipped(klass, reflection, error)
+        ColdStorage.logger.debug do
+          "#{Logging::PREFIX} skipping association #{klass.source_model_name}##{reflection.name} " \
+            "on the archive model: #{error.class}: #{error.message}"
+        end
+      end
+
+      def options_for(reflection, target)
+        options = {
+          class_name: "ColdStorage::Archived::#{const_name_for(target.table_name)}",
+          foreign_key: reflection.foreign_key,
+          inverse_of: false
+        }
+        options[:primary_key] = reflection.options[:primary_key] if reflection.options[:primary_key]
+        options[:as] = reflection.options[:as] if reflection.options[:as]
+        options[:optional] = true if reflection.macro == :belongs_to
+        options
+      end
+
+      def const_name_for(table)
+        name = table.gsub(/[^A-Za-z0-9_]/, '_').camelize
+        name = "T#{name}" unless name.match?(/\A[A-Z]/)
+        name
+      end
+    end
+  end
+end

data/lib/cold_storage/archive_record.rb

@@ -0,0 +1,116 @@
+# frozen_string_literal: true
+
+module ColdStorage
+  # Abstract base of everything that lives in the archive database.
+  #
+  # The connection is established lazily, on first use, so an application that
+  # has not configured the archive database yet still boots.
+  class ArchiveRecord < ActiveRecord::Base
+    self.abstract_class = true
+
+    # Name of the model in the primary database this table belongs to.
+    class_attribute :source_model_name, instance_accessor: false, default: nil
+
+    # The model in the primary database an archived row came from.
+    def self.source_model
+      source_model_name&.constantize
+    end
+
+    def source_model
+      self.class.source_model
+    end
+
+    # Moves this row (and, with `with:`, its archived children) back into the
+    # primary database.
+    #
+    #   Invoice.archived.find(42).restore!(with: :all)
+    #
+    # @return [Integer] number of restored rows
+    def restore!(**options)
+      model = source_model
+      raise NotArchivableError, "#{self.class} has no source model" if model.nil?
+
+      ColdStorage.restore(model, [id], **options)
+    end
+
+    class << self
+      # @return [self]
+      def connect!
+        return self if @connected
+
+        MUTEX.synchronize do
+          next if @connected
+
+          database = ColdStorage.config.archive_database.to_sym
+          assert_configured!(database)
+          connects_to database: { writing: database, reading: database }
+          @connected = true
+        end
+
+        self
+      end
+
+      # Yields a checked out connection to the archive database.
+      # (Named differently from ActiveRecord's own `with_connection` so that
+      # nothing in ActiveRecord ends up calling this override.)
+      def with_archive_connection(&block)
+        connect!
+        connection_pool.with_connection(&block)
+      end
+
+      # Forget the established connection, e.g. after changing the config in a
+      # test. The pool itself is left to ActiveRecord.
+      def reset_connection!
+        MUTEX.synchronize { @connected = false }
+      end
+
+      def database_config
+        env = current_env
+        ActiveRecord::Base.configurations.configs_for(
+          env_name: env,
+          name: ColdStorage.config.archive_database.to_s,
+          include_hidden: true
+        )
+      end
+
+      private
+
+      MUTEX = Mutex.new
+      private_constant :MUTEX
+
+      def current_env
+        if defined?(Rails) && Rails.respond_to?(:env)
+          Rails.env.to_s
+        else
+          ENV['RAILS_ENV'].presence || ENV['RACK_ENV'].presence || 'default_env'
+        end
+      end
+
+      def assert_configured!(database)
+        return if database_config
+
+        available = ActiveRecord::Base.configurations
+                                      .configs_for(env_name: current_env, include_hidden: true)
+                                      .map(&:name)
+        raise ConfigurationError, <<~MESSAGE
+          No "#{database}" database configured for the "#{current_env}" environment.
+          Available entries: #{available.join(', ')}.
+
+          Add one to config/database.yml, for example:
+
+            #{current_env}:
+              primary:
+                <<: *default
+              #{database}:
+                <<: *default
+                database: <%= ENV.fetch('ARCHIVE_POSTGRES_DB') %>
+                database_tasks: false
+
+          `database_tasks: false` keeps `rails db:migrate` from running the
+          application migrations against the archive database; ColdStorage
+          mirrors the schema itself.
+        MESSAGE
+      end
+    end
+  end
+end

data/lib/cold_storage/archiver.rb

@@ -0,0 +1,283 @@
+# frozen_string_literal: true
+
+module ColdStorage
+  # Moves rows of one model from the primary database to the archive database.
+  #
+  # The two databases cannot share a transaction, so every batch is:
+  #   1. read from the primary database,
+  #   2. upserted into the archive database (idempotent, keyed on the primary
+  #      key, so a retry can never duplicate a row),
+  #   3. deleted from the primary database.
+  #
+  # A crash between 2 and 3 leaves the row in both databases; the next run
+  # upserts it again and deletes it. Archiving is therefore at-least-once, and
+  # never loses a row.
+  class Archiver
+    include Logging
+    include RowReader
+
+    # Outcome of one call.
+    Result = Struct.new(:model_name, :archived, :deleted, :skipped, :reason, :dry_run, keyword_init: true) do
+      def skipped?
+        skipped
+      end
+
+      def to_s
+        return "#{model_name}: skipped (#{reason})" if skipped?
+
+        "#{model_name}: archived #{archived}, deleted #{deleted}#{' [dry run]' if dry_run}"
+      end
+    end
+
+    attr_reader :model, :policy, :relation, :result
+
+    # @param model [Class] the ActiveRecord model to archive
+    # @param relation [ActiveRecord::Relation, nil] rows to archive, defaults to
+    #   what the model's policy selects
+    # @param force [Boolean] ignore the `every:` window
+    # @param limit [Integer, nil] stop after this many rows
+    # @param dry_run [Boolean] report what would move without moving it
+    # @param track [Boolean] write a cold_storage_runs row
+    # @param cascade [Boolean, Symbol, Array, Hash] true to follow the policy's
+    #   cascade, false for none, or an explicit association tree
+    # @param delete_after_archive [Boolean, nil] override the policy's setting
+    def initialize(model, relation: nil, force: false, limit: nil,
+                   dry_run: ColdStorage.config.dry_run, batch_size: nil, track: true,
+                   cascade: true, delete_after_archive: nil)
+      @model      = model
+      @policy     = model.respond_to?(:archiving_policy) ? model.archiving_policy : nil
+      @relation   = relation
+      @force      = force
+      @limit      = limit
+      @dry_run    = dry_run
+      @batch_size = batch_size
+      @track      = track && relation.nil?
+      @cascade    = cascade
+      @delete_after_archive = delete_after_archive
+
+      return if @policy || @relation
+
+      raise NotArchivableError,
+            "#{model} is not archivable. Add `archivable ...` to the model, or pass an explicit relation."
+    end
+
+    # @return [Result]
+    def call
+      return skipped(:disabled) unless ColdStorage.config.enabled
+      return skipped(:destroy_only) if sweep_less?
+      return skipped(:not_due) unless due?
+
+      policy&.validate_against_schema!(model)
+      ensure_archive_table!
+
+      run = start_run
+      archived = 0
+      deleted = 0
+
+      each_batch do |rows, ids|
+        archived += copy(rows)
+        deleted  += cascade_and_delete(ids)
+        break if limit_reached?(archived)
+
+        throttle
+      end
+
+      run&.succeed!(archived_count: archived, deleted_count: deleted)
+      finish(archived, deleted)
+    rescue StandardError => e
+      run&.fail!(e, archived_count: archived.to_i, deleted_count: deleted.to_i)
+      raise
+    end
+
+    private
+
+    def due?
+      return true if @force || policy.nil?
+
+      policy.due?(Run.last_success_at(model))
+    end
+
+    # `archivable on_destroy: true` without an age or scope has nothing for a
+    # scheduled run to sweep: it only reacts to deletes.
+    def sweep_less?
+      @relation.nil? && policy && !policy.sweeps?
+    end
+
+    def batch_size
+      @batch_size || policy&.batch_size_value || ColdStorage.config.batch_size
+    end
+
+    def delete_after_archive?
+      return @delete_after_archive unless @delete_after_archive.nil?
+
+      policy.nil? ? ColdStorage.config.delete_after_archive : policy.delete_after_archive?
+    end
+
+    def delete_method
+      policy&.delete_method_value || ColdStorage.config.delete_method
+    end
+
+    def primary_key
+      @primary_key ||= model.primary_key
+    end
+
+    def scope
+      @relation || policy.relation(model)
+    end
+
+    def archive_model
+      @archive_model ||= ArchiveModel.for(model)
+    end
+
+    # Walks the archivable rows in primary key order.
+    #
+    # Reading raw rows (rather than model instances) keeps the values in their
+    # database representation, so model level serializers, default scopes and
+    # callbacks cannot change what gets archived.
+    def each_batch
+      cursor = nil
+      moved = 0
+
+      loop do
+        batch = scope.reorder(primary_key => :asc).limit(remaining(moved) || batch_size)
+        batch = batch.where(model.arel_table[primary_key].gt(cursor)) if cursor
+
+        rows = read(batch)
+        break if rows.empty?
+
+        ids = rows.map { |row| row[primary_key] }
+        cursor = ids.last
+        moved += rows.size
+
+        yield(rows, ids)
+
+        break if rows.size < batch_size
+      end
+    end
+
+    def remaining(moved)
+      return nil unless @limit
+
+      [@limit - moved, 0].max.then { |left| left.zero? ? nil : [left, batch_size].min }
+    end
+
+    def read(batch)
+      sql = batch.select(model.arel_table[Arel.star]).to_sql
+      model.connection_pool.with_connection do |connection|
+        select_rows_as_hashes(connection, sql, "#{model} Archive Load")
+      end
+    end
+
+    def copy(rows)
+      return rows.size if @dry_run
+
+      now = Time.current
+      payload = rows.map do |row|
+        stamped = row.dup
+        stamped[archived_at_column.to_s] = now if archived_at_column
+        stamped
+      end
+
+      archive_model.unscoped.upsert_all(
+        payload,
+        unique_by: primary_key.to_sym,
+        record_timestamps: false,
+        returning: false
+      )
+      rows.size
+    end
+
+    def cascade_and_delete(ids)
+      archive_children(ids)
+      return 0 unless delete_after_archive?
+      return ids.size if @dry_run
+
+      target = model.unscoped.where(primary_key => ids)
+      model.transaction { target.public_send(delete_method) }
+      ids.size
+    end
+
+    # The associations to take along: the policy's cascade by default, or the
+    # explicit tree the caller passed.
+    def cascade_tree
+      return {} if @cascade == false || @cascade.nil?
+      return policy&.cascade || {} if @cascade == true
+
+      AssociationTree.normalize(@cascade)
+    end
+
+    # Children are archived (and deleted) before their parents, so foreign keys
+    # in the primary database stay satisfied at every step.
+    #
+    # A parent that keeps its rows keeps its children's too: `delete_after_archive`
+    # travels down the tree.
+    def archive_children(parent_ids)
+      cascade_tree.each do |association, nested|
+        reflection = Policy.cascade_reflection!(model, association)
+
+        self.class.new(
+          reflection.klass,
+          relation: child_relation(reflection, parent_ids),
+          force: true,
+          dry_run: @dry_run,
+          batch_size: batch_size,
+          track: false,
+          cascade: nested.nil? ? true : nested,
+          delete_after_archive: @delete_after_archive
+        ).call
+      end
+    end
+
+    def child_relation(reflection, parent_ids)
+      scope = reflection.klass.unscoped.where(reflection.foreign_key => parent_ids)
+      scope = scope.where(reflection.type => model.polymorphic_name) if reflection.type
+      scope
+    end
+
+    def ensure_archive_table!
+      table = model.table_name
+      exists = ArchiveRecord.with_archive_connection { |connection| connection.table_exists?(table) }
+      return if exists
+
+      raise SchemaMissingError,
+            "The archive database has no #{table} table yet. " \
+            'Run `rails cold_storage:schema:sync` (or ColdStorage.sync_schema!) first.'
+    end
+
+    def archived_at_column
+      column = ColdStorage.config.archived_at_column
+      return nil if column.nil?
+      return nil if model.column_names.include?(column.to_s)
+
+      column
+    end
+
+    def start_run
+      return nil if @dry_run || !@track
+
+      Run.start!(model)
+    end
+
+    def limit_reached?(archived)
+      @limit && archived >= @limit
+    end
+
+    def throttle
+      sleep(ColdStorage.config.throttle) if ColdStorage.config.throttle.to_f.positive?
+    end
+
+    def skipped(reason)
+      result = Result.new(model_name: model.to_s, archived: 0, deleted: 0, skipped: true,
+                          reason: reason, dry_run: @dry_run)
+      log(result.to_s)
+      result
+    end
+
+    def finish(archived, deleted)
+      result = Result.new(model_name: model.to_s, archived: archived, deleted: deleted,
+                          skipped: false, reason: nil, dry_run: @dry_run)
+      log(result.to_s)
+      result
+    end
+  end
+end

data/lib/cold_storage/association_tree.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module ColdStorage
+  # Normalizes the association trees used by `cascade:` and `with:` into a
+  # plain Hash of name => nested tree.
+  #
+  #   :lines                          => { lines: nil }
+  #   [:lines, :notes]                => { lines: nil, notes: nil }
+  #   [{ lines: [:taxes] }, :notes]   => { lines: [:taxes], notes: nil }
+  #
+  # Nested values are left as given; each level normalizes its own.
+  module AssociationTree
+    module_function
+
+    def normalize(value)
+      case value
+      when nil, false     then {}
+      when Symbol, String then { value.to_sym => nil }
+      when Array          then value.inject({}) { |tree, item| tree.merge(normalize(item)) }
+      when Hash           then value.to_h { |name, nested| [name.to_sym, nested] }
+      else
+        raise ArgumentError, "expected an association name, array or hash, got #{value.inspect}"
+      end
+    end
+  end
+end