active_storage_quota 0.1.0

data/lib/active_storage_quota/audit/result.rb

@@ -0,0 +1,106 @@
+# frozen_string_literal: true
+
+require "set"
+
+module ActiveStorageQuota
+  module Audit
+    # What an audit found.
+    #
+    # Every category keeps a total count and a bounded sample, so auditing an
+    # installation with millions of divergent rows cannot itself become the
+    # problem. Nothing here is ever written back.
+    class Result
+      DEFAULT_SAMPLE_LIMIT = 100
+
+      attr_reader :sample_limit, :consistency, :counts, :samples,
+                  :unsupported_checks, :ignored_record_types
+
+      # +consistency+ is :snapshot when every check ran as same-connection SQL,
+      # and :best_effort when values had to be compared across connections in
+      # Ruby, where concurrent traffic can shift between reads.
+      def initialize(sample_limit: DEFAULT_SAMPLE_LIMIT, consistency: :snapshot)
+        @sample_limit = sample_limit
+        @consistency = consistency
+        @counts = Hash.new(0)
+        @samples = Hash.new { |hash, key| hash[key] = [] }
+        @unsupported_checks = {}
+        @ignored_record_types = Set.new
+      end
+
+      def record(category, **details)
+        category = category.to_sym
+        @counts[category] += 1
+        @samples[category] << Finding.new(category: category, details: details) if
+          @samples[category].size < sample_limit
+
+        self
+      end
+
+      # A check that could not run here, for example comparing charges against
+      # blobs that live on another connection. The audit continues; only this
+      # check is reported as unavailable.
+      def unsupported!(category, reason:)
+        @unsupported_checks[category.to_sym] = reason
+        self
+      end
+
+      def ignore_record_type(record_type)
+        @ignored_record_types << record_type
+        self
+      end
+
+      def clean?
+        findings_count.zero?
+      end
+
+      def findings_count
+        counts.values.sum
+      end
+
+      # The bounded samples, flattened. Use +summary+ for the true totals.
+      def findings
+        samples.values.flatten
+      end
+
+      def sample(category)
+        samples[category.to_sym]
+      end
+
+      def count(category)
+        counts[category.to_sym]
+      end
+
+      # Every category's true total, even where the sample was truncated.
+      def summary
+        counts.dup
+      end
+
+      def truncated?(category)
+        count(category) > sample(category).size
+      end
+
+      def to_h
+        {
+          clean: clean?,
+          consistency: consistency,
+          findings_count: findings_count,
+          summary: summary,
+          samples: samples.transform_values { |list| list.map(&:to_h) },
+          unsupported_checks: unsupported_checks,
+          ignored_record_types: ignored_record_types.to_a.sort
+        }
+      end
+
+      def to_s
+        return "audit clean (#{consistency})" if clean? && unsupported_checks.empty?
+
+        lines = [ "audit found #{findings_count} issue(s) (#{consistency}):" ]
+        summary.sort_by { |_, total| -total }.each do |category, total|
+          lines << "  #{category}: #{total}#{' (sample truncated)' if truncated?(category)}"
+        end
+        unsupported_checks.each { |category, reason| lines << "  #{category}: unavailable (#{reason})" }
+        lines.join("\n")
+      end
+    end
+  end
+end

data/lib/active_storage_quota/audit.rb

@@ -0,0 +1,75 @@
+# frozen_string_literal: true
+
+require "active_storage_quota/audit/finding"
+require "active_storage_quota/audit/result"
+require "active_storage_quota/audit/counter_audit"
+require "active_storage_quota/audit/ledger_audit"
+
+module ActiveStorageQuota
+  # Read-only inspection of quota accounting.
+  #
+  #   result = ActiveStorageQuota.audit
+  #   result.clean?   # => false
+  #   puts result
+  #
+  # The audit performs no writes of any kind. It creates no quota accounts --
+  # reads never do -- repairs nothing, and deletes nothing. Everything it
+  # notices is reported for a human to decide about.
+  #
+  # Three layers of state are compared, and the audit keeps them distinct:
+  #
+  #   ActiveStorage::Attachment   ultimate truth for logical ownership
+  #   ActiveStorageQuota::Charge  canonical accounting ledger
+  #   Account#used_bytes          denormalized aggregate of charges
+  #
+  # so there are two audits internally: a counter audit (account against its
+  # charges) and a ledger audit (charges against attachments and blobs). The
+  # first can always run. The second needs Active Storage loaded, and reports
+  # itself unavailable rather than failing when it is not.
+  #
+  # The audit is unaffected by whether quota enforcement is switched on.
+  module Audit
+    class << self
+      # +sample_limit+ bounds how many examples are kept per category. Totals
+      # are always exact; only the examples are capped.
+      def run(sample_limit: Result::DEFAULT_SAMPLE_LIMIT)
+        result = Result.new(sample_limit: sample_limit, consistency: consistency)
+
+        CounterAudit.new(result: result).run
+
+        if active_storage_available?
+          LedgerAudit.new(result: result, same_connection: same_connection?).run
+        else
+          %i[
+            missing_charge unexpected_charge attachment_count_mismatch
+            charge_byte_size_mismatch charge_missing_blob
+          ].each { |check| result.unsupported!(check, reason: "Active Storage is not loaded") }
+        end
+
+        result
+      end
+
+      private
+        def active_storage_available?
+          defined?(ActiveStorage::Attachment) && defined?(ActiveStorage::Blob)
+        end
+
+        # :snapshot when every comparison is same-connection SQL, so each
+        # individual check sees one consistent view. :best_effort when blobs
+        # live on another connection and have to be compared in Ruby, where
+        # concurrent traffic can shift between reads. Cross-database auditing
+        # still works; it is simply less exact under load.
+        def consistency
+          same_connection? ? :snapshot : :best_effort
+        end
+
+        def same_connection?
+          return false unless active_storage_available?
+
+          ActiveStorageQuota::Record.connection_pool.equal?(ActiveStorage::Record.connection_pool)
+        rescue StandardError
+          false
+        end
+    end
+  end
+end

data/lib/active_storage_quota/backfill/cursor.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+require "json"
+
+module ActiveStorageQuota
+  class Backfill
+    # An opaque, resumable position in a backfill traversal.
+    #
+    # Deliberately not a bare blob id. Carrying the record_type scope lets a
+    # resume be rejected when it would continue under different traversal
+    # semantics, and the version lets the format change later without silently
+    # misreading an old token.
+    #
+    # Encoded as Base64 of JSON through Array#pack, so no gem is needed for it.
+    class Cursor
+      VERSION = 1
+
+      attr_reader :record_type, :last_blob_id
+
+      def initialize(record_type:, last_blob_id:)
+        @record_type = record_type
+        @last_blob_id = last_blob_id.to_s
+      end
+
+      def self.decode(encoded, record_type:)
+        raise ArgumentError, "cursor must be a String, got #{encoded.class}" unless encoded.is_a?(String)
+
+        payload = parse(encoded)
+
+        unless payload["v"] == VERSION
+          raise ArgumentError,
+                "unsupported cursor version #{payload['v'].inspect}; this build reads version #{VERSION}"
+        end
+
+        # Resuming a Contract-scoped traversal under a Document scope would walk
+        # different rows while pretending to continue the same run.
+        if payload["record_type"] != record_type
+          raise ArgumentError,
+                "cursor was produced for record_type #{payload['record_type'].inspect} " \
+                "but this run is scoped to #{record_type.inspect}"
+        end
+
+        last = payload["last_blob_id"]
+        raise ArgumentError, "cursor is missing last_blob_id" unless last.is_a?(String) && !last.empty?
+
+        new(record_type: record_type, last_blob_id: last)
+      end
+
+      def encode
+        [ JSON.generate(v: VERSION, record_type: record_type, last_blob_id: last_blob_id) ].pack("m0")
+      end
+
+      def self.parse(encoded)
+        payload = JSON.parse(encoded.unpack1("m0").to_s)
+        raise ArgumentError unless payload.is_a?(Hash)
+
+        payload
+      rescue JSON::ParserError, ArgumentError, TypeError
+        raise ArgumentError, "cursor is malformed"
+      end
+      private_class_method :parse
+    end
+  end
+end

data/lib/active_storage_quota/backfill/result.rb

@@ -0,0 +1,146 @@
+# frozen_string_literal: true
+
+module ActiveStorageQuota
+  class Backfill
+    # What a backfill run did, or would have done.
+    #
+    # Exact totals, bounded samples, scalars only.
+    class Result
+      DEFAULT_SAMPLE_LIMIT = 100
+
+      COUNTERS = %i[
+        records_scanned attachments_scanned owners_resolved
+        accounts_created charges_created charges_already_present bytes_discovered
+        batches_completed
+      ].freeze
+
+      attr_reader :sample_limit, :strategy, :counts, :samples, :record_types_processed,
+                  :charges_created_precision
+      attr_accessor :next_cursor
+
+      # +charges_created_precision+ says how far to trust charges_created:
+      #
+      #   :planned    a dry run's projection
+      #   :exact      the adapter returned the rows it actually inserted, so
+      #               rows skipped by a conflict are excluded
+      #   :estimated  the adapter cannot report which rows it skipped, so this
+      #               is the number attempted and may overstate what was created
+      #               if live accounting won a race for the same pair
+      def initialize(dry_run:, strategy:, charges_created_precision: :exact,
+                     sample_limit: DEFAULT_SAMPLE_LIMIT)
+        @dry_run = dry_run
+        @strategy = strategy
+        @charges_created_precision = charges_created_precision
+        @sample_limit = sample_limit
+        @totals = Hash.new(0)
+        @counts = Hash.new(0)
+        @samples = Hash.new { |hash, key| hash[key] = [] }
+        @record_types_processed = Set.new
+        @complete = false
+        @next_cursor = nil
+      end
+
+      def dry_run?
+        @dry_run
+      end
+
+      # False when charges_created counts rows attempted rather than rows
+      # confirmed inserted.
+      def charges_created_exact?
+        charges_created_precision == :exact
+      end
+
+      def complete?
+        @complete
+      end
+
+      def complete!
+        @complete = true
+        @next_cursor = nil
+        self
+      end
+
+      COUNTERS.each { |name| define_method(name) { @totals[name] } }
+
+      def increment(counter, by: 1)
+        @totals[counter] += by
+        self
+      end
+
+      def record(category, **details)
+        category = category.to_sym
+        @counts[category] += 1
+        @samples[category] << details if @samples[category].size < sample_limit
+
+        self
+      end
+
+      def processed_record_type(record_type)
+        @record_types_processed << record_type
+        self
+      end
+
+      def count(category)
+        @counts[category.to_sym]
+      end
+
+      def sample(category)
+        @samples[category.to_sym]
+      end
+
+      # Backfill writes charges but never counters, so any charge it created
+      # leaves used_bytes understating the ledger until counter reconciliation
+      # runs. That is the one sanctioned exception to the runtime invariant, and
+      # it is only safe while enforcement is off.
+      def counters_require_reconciliation?
+        !dry_run? && charges_created.positive?
+      end
+
+      def enforcement_active_owners
+        count(:enforcement_active_owner)
+      end
+
+      def errors
+        count(:errors)
+      end
+
+      def to_h
+        {
+          dry_run: dry_run?,
+          strategy: strategy,
+          complete: complete?,
+          next_cursor: next_cursor,
+          record_types_processed: record_types_processed.to_a.sort,
+          charges_created_precision: charges_created_precision,
+          totals: COUNTERS.to_h { |name| [ name, @totals[name] ] },
+          categories: @counts.dup,
+          counters_require_reconciliation: counters_require_reconciliation?,
+          samples: @samples.transform_values(&:dup)
+        }
+      end
+
+      def to_s
+        headline = dry_run? ? "backfill plan" : "backfill"
+        lines = [
+          "#{headline} (#{strategy}): #{batches_completed} batch(es), " \
+            "#{attachments_scanned} attachment(s), +#{accounts_created} account(s), " \
+            "+#{charges_created} charge(s)#{charge_count_note}, #{bytes_discovered} byte(s) discovered"
+        ]
+        @counts.sort_by { |_, total| -total }.each { |category, total| lines << "  #{category}: #{total}" }
+        lines << "  complete: #{complete?}"
+        lines << "  next cursor: #{next_cursor}" if next_cursor
+        lines << "  run reconcile_counters! before enabling enforcement" if counters_require_reconciliation?
+        lines.join("\n")
+      end
+
+      private
+        def charge_count_note
+          case charges_created_precision
+          when :estimated then " (attempted; this adapter cannot report skipped rows)"
+          when :planned then " (planned)"
+          else ""
+          end
+        end
+    end
+  end
+end

data/lib/active_storage_quota/backfill.rb

@@ -0,0 +1,310 @@
+# frozen_string_literal: true
+
+require "active_storage_quota/backfill/cursor"
+require "active_storage_quota/backfill/result"
+
+module ActiveStorageQuota
+  # Builds the initial charge ledger for an application that installed this gem
+  # after it already had Active Storage attachments.
+  #
+  #   ActiveStorageQuota::Backfill.run(dry_run: true)
+  #   ActiveStorageQuota::Backfill.run(max_batches: 100)
+  #
+  # Backfill is an operational subsystem, not an application operation, which is
+  # why it is not reachable from the top-level module.
+  #
+  # It is strictly ADDITIVE. It inserts quota accounts and charges and does
+  # nothing else: it never updates or deletes a charge, never touches
+  # +used_bytes+, never removes an account and never writes an Active Storage
+  # row. Contradictions it finds are reported for audit and reconciliation.
+  #
+  # Because it does not write +used_bytes+, a completed backfill leaves counters
+  # understating the ledger until +reconcile_counters!+ runs. That is the single
+  # sanctioned exception to the runtime invariant
+  #
+  #   account.used_bytes == account.charges.sum(:byte_size)
+  #
+  # and it is only safe while quota limits resolve to nil. See the README
+  # section "Adopting active_storage_quota in an existing application".
+  class Backfill
+    DEFAULT_BATCH_SIZE = 500
+
+    class << self
+      def run(record_type: nil, cursor: nil, batch_size: DEFAULT_BATCH_SIZE,
+              max_batches: nil, dry_run: false, allow_active_enforcement: false,
+              sample_limit: Result::DEFAULT_SAMPLE_LIMIT)
+        new(
+          record_type: validate_record_type(record_type),
+          batch_size: validate_batch_size(batch_size),
+          max_batches: validate_max_batches(max_batches),
+          cursor: cursor,
+          dry_run: dry_run,
+          allow_active_enforcement: allow_active_enforcement,
+          sample_limit: sample_limit
+        ).run
+      end
+
+      private
+        def validate_batch_size(value)
+          return value if value.is_a?(Integer) && value.positive?
+
+          raise ArgumentError, "batch_size must be a positive Integer, got #{value.inspect}"
+        end
+
+        def validate_max_batches(value)
+          return value if value.nil? || (value.is_a?(Integer) && value.positive?)
+
+          raise ArgumentError, "max_batches must be nil or a positive Integer, got #{value.inspect}"
+        end
+
+        def validate_record_type(value)
+          return value if value.nil? || (value.is_a?(String) && !value.strip.empty?)
+
+          raise ArgumentError, "record_type must be nil or a non-empty String, got #{value.inspect}"
+        end
+    end
+
+    def initialize(record_type:, batch_size:, max_batches:, cursor:, dry_run:,
+                   allow_active_enforcement:, sample_limit:)
+      @record_type = record_type
+      @batch_size = batch_size
+      @max_batches = max_batches
+      @dry_run = dry_run
+      @allow_active_enforcement = allow_active_enforcement
+      @result = Result.new(dry_run: dry_run, strategy: strategy, sample_limit: sample_limit,
+                           charges_created_precision: charges_created_precision(dry_run))
+      @resolver = OwnerResolver.new(result: @result)
+      @cursor = cursor && Cursor.decode(cursor, record_type: record_type)
+      @limits_seen = {}
+    end
+
+    def run
+      last_blob_id = @cursor&.last_blob_id
+      batches = 0
+
+      loop do
+        blob_ids = next_blob_ids(last_blob_id)
+        break result.complete! if blob_ids.empty?
+
+        process(blob_ids)
+
+        last_blob_id = blob_ids.last.to_s
+        batches += 1
+        result.increment(:batches_completed)
+
+        if @max_batches && batches >= @max_batches
+          result.next_cursor = Cursor.new(record_type: @record_type, last_blob_id: last_blob_id).encode
+          break
+        end
+      end
+
+      result
+    end
+
+    private
+      attr_reader :result
+
+      # Keyset over distinct blob ids. Ordering only has to be stable and total,
+      # which holds for bigint, uuid and text keys alike; nothing here assumes
+      # ids are numeric or chronological, and MAX(id) is never consulted.
+      #
+      # Batching by blob keeps every (account, blob) group inside one batch, so
+      # counts are exact.
+      def next_blob_ids(last_blob_id)
+        scope = ActiveStorage::Attachment.order(:blob_id).limit(@batch_size)
+        scope = scope.where(record_type: @record_type) if @record_type
+        if last_blob_id
+          scope = scope.where(ActiveStorage::Attachment.arel_table[:blob_id].gt(last_blob_id))
+        end
+
+        scope.distinct.pluck(:blob_id)
+      end
+
+      def process(blob_ids)
+        owner_pairs = @resolver.owner_pairs_for_blobs(blob_ids)
+        result.increment(:attachments_scanned, by: owner_pairs.values.sum)
+        return if owner_pairs.empty?
+
+        owner_keys = owner_pairs.keys.map { |type, id, _blob| [ type, id ] }.uniq
+        result.increment(:owners_resolved, by: owner_keys.size)
+        record_types_from(owner_pairs)
+
+        # Checked before anything in this batch is written, so a run that would
+        # enforce against incomplete counters stops before it starts.
+        guard_enforcement!(owner_keys)
+
+        accounts = ensure_accounts(owner_keys)
+        sizes = blob_sizes(blob_ids)
+        existing = existing_charges(accounts.values.compact, blob_ids)
+
+        insert_charges(build_rows(owner_pairs, accounts, sizes, existing))
+      end
+
+      def record_types_from(owner_pairs)
+        result.processed_record_type(@record_type) if @record_type
+        result.increment(:records_scanned, by: owner_pairs.size)
+      end
+
+      def build_rows(owner_pairs, accounts, sizes, existing)
+        owner_pairs.filter_map do |(owner_type, owner_id, blob_id), count|
+          account_id = accounts[[ owner_type, owner_id ]]
+          byte_size = sizes[blob_id]
+
+          # Never invent a size. An attachment pointing at a missing blob is
+          # malformed data to report, not to guess at.
+          if byte_size.nil?
+            result.record(:missing_blob, owner_type: owner_type, owner_id: owner_id, blob_id: blob_id)
+            next
+          end
+
+          if account_id && (charge = existing[[ account_id, blob_id ]])
+            note_existing(charge, account_id, blob_id, byte_size)
+            next
+          end
+
+          next if account_id.nil? && !@dry_run
+
+          { quota_account_id: account_id, blob_id: blob_id, byte_size: byte_size,
+            attachment_count: count }
+        end
+      end
+
+      # An existing charge is left completely alone -- count, byte size, owner.
+      # Repairing those is reconciliation's job, and doing it here would mean a
+      # second concurrency model for the same problem.
+      def note_existing(charge, account_id, blob_id, byte_size)
+        result.increment(:charges_already_present)
+        return if charge == byte_size
+
+        result.record(:report_only_mismatch, quota_account_id: account_id, blob_id: blob_id,
+                                             charge_byte_size: charge, blob_byte_size: byte_size)
+      end
+
+      def insert_charges(rows)
+        return if rows.empty?
+
+        result.increment(:bytes_discovered, by: rows.sum { |row| row[:byte_size] })
+
+        if @dry_run
+          result.increment(:charges_created, by: rows.size)
+          return
+        end
+
+        # Conflict means skip, never upsert: a live callback's delta must never
+        # be overwritten with a historical absolute count.
+        inserted = Charge.insert_all(rows, unique_by: %i[quota_account_id blob_id],
+                                           record_timestamps: true)
+        result.increment(:charges_created, by: inserted_count(inserted, rows))
+      end
+
+      # insert_all only asks for RETURNING when the adapter supports it
+      # (ActiveRecord::InsertAll sets @returning from supports_insert_returning?).
+      # Where it does, the returned rows are exactly the ones inserted, so a row
+      # skipped by a conflict is correctly excluded. Where it does not, the
+      # result carries no rows at all -- reading its length would report zero --
+      # so the count falls back to rows attempted and the result says so rather
+      # than presenting an estimate as a fact.
+      def inserted_count(inserted, rows)
+        return rows.size unless result.charges_created_exact?
+
+        inserted.length
+      end
+
+      def charges_created_precision(dry_run)
+        return :planned if dry_run
+
+        insert_returning_supported? ? :exact : :estimated
+      end
+
+      def insert_returning_supported?
+        Charge.connection_pool.with_connection(&:supports_insert_returning?)
+      rescue StandardError
+        false
+      end
+
+      def ensure_accounts(owner_keys)
+        existing = @resolver.send(:accounts_by_owner, owner_keys)
+        missing = owner_keys - existing.keys
+        return existing if missing.empty?
+
+        if @dry_run
+          result.increment(:accounts_created, by: missing.size)
+          return existing
+        end
+
+        Account.insert_all(
+          missing.map do |owner_type, owner_id|
+            { owner_type: owner_type, owner_id: owner_id,
+              scope_name: ActiveStorageQuota::DEFAULT_SCOPE, used_bytes: 0, reserved_bytes: 0 }
+          end,
+          unique_by: %i[owner_type owner_id scope_name],
+          record_timestamps: true
+        )
+
+        refreshed = @resolver.send(:accounts_by_owner, owner_keys)
+        result.increment(:accounts_created, by: (refreshed.keys - existing.keys).size)
+        refreshed
+      end
+
+      # Backfill leaves used_bytes understating the ledger until counter
+      # reconciliation runs, so enforcing against it would admit uploads that
+      # should have been refused.
+      def guard_enforcement!(owner_keys)
+        active = owner_keys.select { |key| enforcement_active?(key) }
+        return if active.empty?
+
+        active.each do |owner_type, owner_id|
+          result.record(:enforcement_active_owner, owner_type: owner_type, owner_id: owner_id)
+        end
+
+        return if @dry_run || @allow_active_enforcement
+
+        raise EnforcementActive.new(owners: active)
+      end
+
+      def enforcement_active?(owner_key)
+        @limits_seen.fetch(owner_key) { @limits_seen[owner_key] = resolve_limit(owner_key) }
+      end
+
+      def resolve_limit(owner_key)
+        owner_type, owner_id = owner_key
+        klass = owner_type.safe_constantize
+        return false if klass.nil?
+
+        owner = klass.find_by(klass.primary_key => owner_id)
+        return false if owner.nil? || !owner.respond_to?(:storage_limit)
+
+        !owner.storage_limit.nil?
+      rescue StandardError
+        false
+      end
+
+      def blob_sizes(blob_ids)
+        ActiveStorage::Blob.where(id: blob_ids).pluck(:id, :byte_size)
+                           .each_with_object({}) { |(id, size), map| map[id.to_s] = size }
+      end
+
+      def existing_charges(account_ids, blob_ids)
+        return {} if account_ids.empty?
+
+        Charge.where(quota_account_id: account_ids, blob_id: blob_ids)
+              .pluck(:quota_account_id, :blob_id, :byte_size)
+              .each_with_object({}) { |(account_id, blob_id, size), map| map[[ account_id, blob_id.to_s ]] = size }
+      end
+
+      # Reads come from Active Storage, writes go to the quota tables. Neither
+      # depends on a cross-database join, so a separate Active Storage database
+      # costs the strategy label and nothing else.
+      def strategy
+        self.class.send(:same_connection?) ? :same_connection : :cross_connection
+      end
+
+      class << self
+        def same_connection?
+          ActiveStorageQuota::Record.connection_pool.equal?(ActiveStorage::Record.connection_pool)
+        rescue StandardError
+          false
+        end
+      end
+  end
+end