partition_gardener 0.3.0

data/lib/partition_gardener/templates.rb

@@ -0,0 +1,373 @@
+module PartitionGardener
+  module Templates
+    module_function
+
+    def normalize(config)
+      config = config.dup
+      config[:layout] ||= :sliding_window
+      config[:bucket] ||= :month if date_layout?(config[:layout])
+      config
+    end
+
+    def expand(config)
+      normalized = normalize(config)
+      return [normalized] unless normalized[:layout] == :composite
+
+      configs = []
+      parent_mode = normalized.fetch(:parent_mode, :list)
+
+      case parent_mode
+      when :list
+        if normalized[:list_branches]
+          configs << list_split(
+            table_name: normalized[:table_name],
+            branches: normalized[:list_branches],
+            conflict_key: normalized[:conflict_key],
+            partition_key_column: normalized.fetch(:partition_key_column, normalized[:discriminator_column])
+          )
+        end
+      when :range
+        configs << range_parent_config_for(normalized)
+      end
+
+      normalized.fetch(:branches, []).each do |branch|
+        configs << branch_config_for(normalized, branch)
+      end
+
+      configs
+    end
+
+    def sliding_window_monthly(table_name:, partition_key_column:, conflict_key:, active_months: 12, split_row_threshold: nil, **options)
+      sliding_window_for_bucket(
+        :month,
+        table_name: table_name,
+        partition_key_column: partition_key_column,
+        conflict_key: conflict_key,
+        active_months: active_months,
+        split_row_threshold: split_row_threshold,
+        **options
+      )
+    end
+
+    def sliding_window_daily(table_name:, partition_key_column:, conflict_key:, active_days: 90, split_row_threshold: nil, **options)
+      sliding_window_for_bucket(
+        :day,
+        table_name: table_name,
+        partition_key_column: partition_key_column,
+        conflict_key: conflict_key,
+        active_days: active_days,
+        split_row_threshold: split_row_threshold,
+        **options
+      )
+    end
+
+    def sliding_window_weekly(table_name:, partition_key_column:, conflict_key:, active_weeks: 52, split_row_threshold: nil, **options)
+      sliding_window_for_bucket(
+        :week,
+        table_name: table_name,
+        partition_key_column: partition_key_column,
+        conflict_key: conflict_key,
+        active_weeks: active_weeks,
+        split_row_threshold: split_row_threshold,
+        **options
+      )
+    end
+
+    def sliding_window_quarterly(table_name:, partition_key_column:, conflict_key:, active_quarters: 8, split_row_threshold: nil, **options)
+      sliding_window_for_bucket(
+        :quarter,
+        table_name: table_name,
+        partition_key_column: partition_key_column,
+        conflict_key: conflict_key,
+        active_quarters: active_quarters,
+        split_row_threshold: split_row_threshold,
+        **options
+      )
+    end
+
+    def rolling_current_monthly(table_name:, partition_key_column:, conflict_key:, active_months: 12, **options)
+      sliding_window_monthly(
+        table_name: table_name,
+        partition_key_column: partition_key_column,
+        conflict_key: conflict_key,
+        active_months: active_months,
+        layout: :rolling_current,
+        split_row_threshold: Float::INFINITY,
+        **options
+      )
+    end
+
+    def premake_monthly(table_name:, partition_key_column:, conflict_key:, premake_months: 3, split_row_threshold: nil, **options)
+      normalize({
+        table_name: table_name,
+        layout: :premake_monthly,
+        bucket: :month,
+        partition_key_column: partition_key_column,
+        conflict_key: conflict_key,
+        premake_months: premake_months,
+        maintenance_backend: options.fetch(:maintenance_backend, :gardener),
+        partition_name_format: options.fetch(:partition_name_format) {
+          ->(identifier) { DateBucket.partition_name(table_name, identifier, :month) }
+        },
+        partition_definition: options.fetch(:partition_definition) {
+          ->(date) { DateBucket.partition_definition_clause(date, :month) }
+        },
+        extract_partition_identifier: options.fetch(:extract_partition_identifier) {
+          ->(date_value) { DateBucket.beginning_of_bucket(date_value, :month) }
+        }
+      }.merge(options.except(:partition_name_format, :partition_definition, :extract_partition_identifier)).tap do |config|
+        config[:split_row_threshold] = split_row_threshold if split_row_threshold
+      end)
+    end
+
+    def calendar_year(table_name:, partition_key_column:, conflict_key:, active_years: 2, split_row_threshold: nil, **options)
+      normalize({
+        table_name: table_name,
+        layout: :calendar_year,
+        bucket: :year,
+        partition_key_column: partition_key_column,
+        conflict_key: conflict_key,
+        active_years: active_years,
+        partition_name_format: options.fetch(:partition_name_format) {
+          ->(identifier) { DateBucket.partition_name(table_name, identifier, :year) }
+        },
+        partition_definition: options.fetch(:partition_definition) {
+          ->(date) { DateBucket.partition_definition_clause(date, :year) }
+        },
+        extract_partition_identifier: options.fetch(:extract_partition_identifier) {
+          ->(date_value) { DateBucket.beginning_of_bucket(date_value, :year) }
+        }
+      }.merge(options.except(:partition_name_format, :partition_definition, :extract_partition_identifier)).tap do |config|
+        config[:split_row_threshold] = split_row_threshold if split_row_threshold
+      end)
+    end
+
+    def integer_window(table_name:, partition_key_column:, conflict_key:, active_id_lo: 0, active_id_width: nil, current_band_size: nil, split_row_threshold: nil, **options)
+      normalize({
+        table_name: table_name,
+        layout: :integer_window,
+        partition_key_column: partition_key_column,
+        conflict_key: conflict_key,
+        active_id_lo: active_id_lo,
+        active_id_width: active_id_width || Strategy::IntegerRange::DEFAULT_ACTIVE_ID_WIDTH,
+        current_band_size: current_band_size || Strategy::IntegerRange::DEFAULT_CURRENT_BAND_SIZE,
+        archive_band_size: options.fetch(:archive_band_size, Strategy::IntegerRange::DEFAULT_ARCHIVE_BAND_SIZE)
+      }.merge(options).tap do |config|
+        config[:split_row_threshold] = split_row_threshold if split_row_threshold
+      end)
+    end
+
+    def hash_branches(table_name:, partition_key_column:, conflict_key:, hash_modulus: nil, split_row_threshold: nil, **options)
+      normalize({
+        table_name: table_name,
+        layout: :hash_branches,
+        partition_key_column: partition_key_column,
+        conflict_key: conflict_key,
+        hash_modulus: hash_modulus || Strategy::HashBranches::DEFAULT_MODULUS
+      }.merge(options).tap do |config|
+        config[:split_row_threshold] = split_row_threshold if split_row_threshold
+      end)
+    end
+
+    def list_split(table_name:, branches:, conflict_key:, partition_key_column: nil, **options)
+      discriminator_column = partition_key_column || branches.first&.fetch(:discriminator_column, nil)
+      normalized_branches = Predicate.normalize_branches!(
+        branches,
+        discriminator_column: discriminator_column
+      )
+
+      normalize({
+        table_name: table_name,
+        layout: :list_split,
+        branches: normalized_branches,
+        conflict_key: conflict_key,
+        partition_key_column: partition_key_column || discriminator_column || "id"
+      }.merge(options))
+    end
+
+    def composite_list_hash(parent_table:, discriminator_column:, branches:, conflict_key:, **options)
+      composite_with_branches(
+        parent_table: parent_table,
+        discriminator_column: discriminator_column,
+        branches: branches,
+        conflict_key: conflict_key,
+        parent_mode: :list,
+        default_child_layout: :hash_branches,
+        **options
+      )
+    end
+
+    def composite_list_range(parent_table:, discriminator_column:, branches:, conflict_key:, bucket: :month, **options)
+      composite_with_branches(
+        parent_table: parent_table,
+        discriminator_column: discriminator_column,
+        branches: branches,
+        conflict_key: conflict_key,
+        parent_mode: :list,
+        default_child_layout: :sliding_window,
+        default_child_bucket: bucket,
+        **options
+      )
+    end
+
+    def list_range(parent_table:, discriminator_column:, branches:, conflict_key:, **options)
+      composite_list_range(
+        parent_table: parent_table,
+        discriminator_column: discriminator_column,
+        branches: branches,
+        conflict_key: conflict_key,
+        **options
+      )
+    end
+
+    def composite_range_hash(parent_table:, partition_key_column:, branches:, conflict_key:, bucket: :month, **options)
+      composite_with_branches(
+        parent_table: parent_table,
+        partition_key_column: partition_key_column,
+        branches: branches,
+        conflict_key: conflict_key,
+        parent_mode: :range,
+        parent_bucket: bucket,
+        default_child_layout: :hash_branches,
+        **options
+      )
+    end
+
+    def composite_range_list(parent_table:, partition_key_column:, branches:, conflict_key:, bucket: :month, **options)
+      composite_with_branches(
+        parent_table: parent_table,
+        partition_key_column: partition_key_column,
+        branches: branches,
+        conflict_key: conflict_key,
+        parent_mode: :range,
+        parent_bucket: bucket,
+        default_child_layout: :list_split,
+        **options
+      )
+    end
+
+    def branch_config_for(parent_config, branch)
+      branch_table_name = "#{parent_config[:table_name]}_#{branch.fetch(:name)}"
+      child = branch.dup
+      child.delete(:name)
+      child[:table_name] = branch_table_name
+      child[:parent_table_name] = parent_config[:table_name]
+      child[:layout] ||= parent_config.fetch(:default_child_layout, :hash_branches)
+
+      if child[:layout] == :sliding_window
+        child[:bucket] ||= parent_config.fetch(:default_child_bucket, :month)
+        child[:partition_key_column] ||= branch.fetch(:partition_key_column)
+        return sliding_window_for_bucket(
+          child[:bucket],
+          table_name: branch_table_name,
+          partition_key_column: child[:partition_key_column],
+          conflict_key: parent_config[:conflict_key],
+          parent_table_name: parent_config[:table_name],
+          split_row_threshold: child[:split_row_threshold],
+          active_months: child[:active_months],
+          active_days: child[:active_days],
+          active_weeks: child[:active_weeks],
+          active_quarters: child[:active_quarters]
+        )
+      end
+
+      if child[:layout] == :list_split
+        child[:branches] = branch.fetch(:branches)
+        child[:conflict_key] = parent_config[:conflict_key]
+        return list_split(
+          table_name: branch_table_name,
+          branches: child[:branches],
+          conflict_key: child[:conflict_key],
+          partition_key_column: child[:partition_key_column],
+          parent_table_name: parent_config[:table_name]
+        )
+      end
+
+      normalize(child)
+    end
+
+    def date_layout?(layout)
+      layout == :sliding_window || layout == :calendar_year || layout == :premake_monthly || layout == :rolling_current
+    end
+
+    def sliding_window_for_bucket(bucket, table_name:, partition_key_column:, conflict_key:, split_row_threshold: nil, layout: :sliding_window, **options)
+      active_key = DateBucket.active_key(bucket)
+      active_value = options[active_key] || DateBucket.default_active_span(bucket)
+
+      normalize({
+        :table_name => table_name,
+        :layout => layout,
+        :bucket => bucket,
+        :partition_key_column => partition_key_column,
+        :conflict_key => conflict_key,
+        active_key => active_value,
+        :partition_name_format => options.fetch(:partition_name_format) {
+          ->(identifier) { DateBucket.partition_name(table_name, identifier, bucket) }
+        },
+        :partition_definition => options.fetch(:partition_definition) {
+          ->(date) { DateBucket.partition_definition_clause(date, bucket) }
+        },
+        :extract_partition_identifier => options.fetch(:extract_partition_identifier) {
+          ->(date_value) { DateBucket.beginning_of_bucket(date_value, bucket) }
+        }
+      }.merge(options.except(:partition_name_format, :partition_definition, :extract_partition_identifier, active_key)).tap do |config|
+        config[:split_row_threshold] = split_row_threshold if split_row_threshold
+      end)
+    end
+
+    def composite_with_branches(parent_table:, branches:, conflict_key:, parent_mode:, default_child_layout:, discriminator_column: nil, partition_key_column: nil, parent_bucket: :month, default_child_bucket: :month, **options)
+      list_branch_entries = Predicate.list_branch_entries(branches, discriminator_column: discriminator_column) if parent_mode == :list
+
+      child_branches = branches.map do |branch|
+        entry = {
+          name: branch.fetch(:name),
+          layout: branch.fetch(:layout, default_child_layout)
+        }
+        entry[:partition_key_column] = branch[:partition_key_column] if branch[:partition_key_column]
+        entry[:hash_modulus] = branch.fetch(:hash_modulus, Strategy::HashBranches::DEFAULT_MODULUS) if entry[:layout] == :hash_branches
+        entry[:bucket] = branch.fetch(:bucket, default_child_bucket) if entry[:layout] == :sliding_window
+        entry[:active_months] = branch[:active_months] if branch[:active_months]
+        entry[:active_days] = branch[:active_days] if branch[:active_days]
+        entry[:active_weeks] = branch[:active_weeks] if branch[:active_weeks]
+        entry[:active_quarters] = branch[:active_quarters] if branch[:active_quarters]
+        entry[:split_row_threshold] = branch[:split_row_threshold] if branch[:split_row_threshold]
+        entry[:branches] = branch[:branches] if branch[:branches]
+        entry
+      end
+
+      normalize({
+        table_name: parent_table,
+        layout: :composite,
+        parent_mode: parent_mode,
+        bucket: parent_bucket,
+        discriminator_column: discriminator_column,
+        partition_key_column: partition_key_column,
+        conflict_key: conflict_key,
+        default_child_layout: default_child_layout,
+        default_child_bucket: default_child_bucket,
+        list_branches: (parent_mode == :list) ? list_branch_entries : nil,
+        branches: child_branches
+      }.compact.merge(options))
+    end
+
+    def range_parent_config_for(parent_config)
+      bucket = parent_config.fetch(:bucket, :month)
+      base = {
+        table_name: parent_config[:table_name],
+        partition_key_column: parent_config.fetch(:partition_key_column),
+        conflict_key: parent_config[:conflict_key]
+      }
+      passthrough = parent_config.slice(
+        :active_months, :active_days, :active_weeks, :active_quarters, :active_years,
+        :split_row_threshold, :maintenance_backend, :retention_months, :move_batch_size,
+        :statement_timeout, :incremental_rebalance, :analyze_after_rebalance
+      )
+
+      if bucket == :year
+        calendar_year(**base, **passthrough)
+      else
+        sliding_window_for_bucket(bucket, **base, **passthrough)
+      end
+    end
+  end
+end

data/lib/partition_gardener/unmoved_rows_remaining.rb

@@ -0,0 +1,15 @@
+module PartitionGardener
+  class UnmovedRowsRemaining < StandardError
+    attr_reader :source_partition_name, :batch_size, :last_cursor
+
+    def initialize(source_partition_name:, batch_size:, last_cursor:)
+      @source_partition_name = source_partition_name
+      @batch_size = batch_size
+      @last_cursor = last_cursor
+      super(
+        "Stopped moving rows from #{source_partition_name}: " \
+        "batch selected #{batch_size} row(s) but none could be removed"
+      )
+    end
+  end
+end

data/lib/partition_gardener/version.rb

@@ -0,0 +1,3 @@
+module PartitionGardener
+  VERSION = "0.3.0"
+end

data/lib/partition_gardener.rb

@@ -0,0 +1,215 @@
+require "date"
+
+require_relative "partition_gardener/version"
+require_relative "partition_gardener/date_calendar"
+require_relative "partition_gardener/date_bucket"
+require_relative "partition_gardener/stdlib_extensions"
+require "pg"
+require_relative "partition_gardener/pg_connection"
+require_relative "partition_gardener/blank"
+require_relative "partition_gardener/predicate"
+require_relative "partition_gardener/memory_run_record_store"
+require_relative "partition_gardener/sql_run_record_store"
+require_relative "partition_gardener/configuration"
+require_relative "partition_gardener/naming"
+require_relative "partition_gardener/plan"
+require_relative "partition_gardener/connection"
+require_relative "partition_gardener/executor"
+require_relative "partition_gardener/default_partition"
+require_relative "partition_gardener/strategy/requires_default_partition"
+require_relative "partition_gardener/strategy/cursor_columns"
+require_relative "partition_gardener/strategy/date_range"
+require_relative "partition_gardener/strategy/integer_range"
+require_relative "partition_gardener/strategy/hash_branches"
+require_relative "partition_gardener/strategy/list_split"
+require_relative "partition_gardener/strategy/composite"
+require_relative "partition_gardener/strategy"
+require_relative "partition_gardener/layout/three_area"
+require_relative "partition_gardener/layout/zone_segments"
+require_relative "partition_gardener/layout/sliding_window"
+require_relative "partition_gardener/layout/calendar_year"
+require_relative "partition_gardener/layout/integer_window"
+require_relative "partition_gardener/planner"
+require_relative "partition_gardener/plan_diff"
+require_relative "partition_gardener/gap_detection"
+require_relative "partition_gardener/plan_report"
+require_relative "partition_gardener/run_record"
+require_relative "partition_gardener/plan_applier"
+require_relative "partition_gardener/hash_routing"
+require_relative "partition_gardener/maintenance_backend"
+require_relative "partition_gardener/run_metrics"
+require_relative "partition_gardener/date_range_maintenance"
+require_relative "partition_gardener/premake_monthly_maintenance"
+require_relative "partition_gardener/templates"
+require_relative "partition_gardener/advisory_lock"
+require_relative "partition_gardener/lock_not_acquired"
+require_relative "partition_gardener/missing_conflict_index"
+require_relative "partition_gardener/orphaned_rebalance_staging"
+require_relative "partition_gardener/unmoved_rows_remaining"
+require_relative "partition_gardener/run_failed"
+require_relative "partition_gardener/audit"
+require_relative "partition_gardener/archive_retention"
+require_relative "partition_gardener/registry"
+require_relative "partition_gardener/config_document"
+require_relative "partition_gardener/cli"
+require_relative "partition_gardener/migration/hot_switch_concern"
+
+module PartitionGardener
+  MOVE_BATCH_SIZE = 10_000
+  FUTURE_MONTH_PARTITION_ROW_THRESHOLD = 100_000
+  DEFAULT_STATEMENT_TIMEOUT = 300
+
+  class << self
+    def run!(statement_timeout: DEFAULT_STATEMENT_TIMEOUT, job_class_name: "PartitionGardener", continue_on_error: configuration.continue_on_error, dry_run: false, table_name: nil)
+      return plan_all if dry_run
+
+      errors = []
+      table_results = []
+      filter_table = table_name
+      configs = filter_table ? Registry.configs_for_table(filter_table) : Registry.expanded_table_configs
+
+      configs.each do |config|
+        table_results << run_single_table!(
+          config,
+          statement_timeout: statement_timeout,
+          job_class_name: job_class_name,
+          continue_on_error: continue_on_error,
+          errors: errors
+        )
+      end
+
+      summary = RunSummary.new(tables: table_results, errors: errors)
+      raise RunFailed, errors if errors.any?
+
+      summary
+    end
+
+    def audit(table_name)
+      configs = Registry.configs_for_table(table_name)
+      raise ArgumentError, "no config for #{table_name}" if configs.empty?
+
+      if configs.one?
+        Audit.call(table_name, config: configs.first)
+      else
+        configs.map { |config| Audit.call(config[:table_name], config: config) }
+      end
+    end
+
+    def plan(table_name: nil, config: nil)
+      if config
+        return PlanReport.build(config)
+      end
+
+      raise ArgumentError, "table_name or config is required" unless table_name
+
+      configs = Registry.configs_for_table(table_name)
+      raise ArgumentError, "no config for #{table_name}" if configs.empty?
+
+      if configs.one?
+        PlanReport.build(configs.first)
+      else
+        configs.map { |entry| PlanReport.build(entry) }
+      end
+    end
+
+    def maintenance_for(config, job_class_name: "PartitionGardener")
+      case config.fetch(:layout)
+      when :composite
+        CompositeMaintenance.new(config, job_class_name: job_class_name)
+      when :premake_monthly
+        PremakeMonthlyMaintenance.new(config, job_class_name: job_class_name)
+      else
+        ThreeAreaMaintenance.new(config, job_class_name: job_class_name)
+      end
+    end
+
+    def emit_run_metadata(metrics, job_class_name:)
+      configuration.notify(
+        "[PartitionGardener] Finished #{metrics.table_name}",
+        context: {
+          table_name: metrics.table_name,
+          job: job_class_name,
+          action: "run_metadata",
+          run_metadata: metrics.to_h
+        }
+      )
+    end
+
+    def plan_all
+      reports = []
+      Registry.each_table_config do |config|
+        next unless Connection.table_is_partitioned?(config[:table_name])
+
+        reports << PlanReport.build(config).to_h
+      end
+      reports
+    end
+
+    def plan_all_tables
+      Registry.each_table_config.map { |config| PlanReport.build(config) }
+    end
+
+    def run_single_table!(config, statement_timeout:, job_class_name:, continue_on_error:, errors:)
+      table_name = config[:table_name]
+      table_timeout = config.fetch(:statement_timeout, statement_timeout)
+      metrics = RunMetrics.new(table_name)
+
+      configuration.with_statement_timeout(table_timeout) do
+        Connection.clear_attached_partitions_cache!
+
+        unless Connection.table_is_partitioned?(table_name)
+          metrics.mark_skipped!("not_partitioned")
+          return metrics
+        end
+
+        if MaintenanceBackend.skipped?(config)
+          metrics.mark_skipped!("maintenance_backend_pg_partman")
+          return metrics
+        end
+
+        configuration.current_run_metrics = metrics
+
+        AdvisoryLock.with_table_lock(table_name) do
+          DefaultPartition.ensure!(config)
+          maintenance_for(config, job_class_name: job_class_name).run!
+        end
+
+        metrics.finish!
+        emit_run_metadata(metrics, job_class_name: job_class_name)
+        metrics
+      rescue LockNotAcquired => error
+        metrics.mark_skipped!("lock_not_acquired")
+        configuration.notify(
+          "[PartitionGardener] skipped #{table_name}: #{error.message}",
+          context: {
+            table_name: table_name,
+            job: job_class_name,
+            action: "lock"
+          }
+        )
+        metrics
+      rescue => error
+        metrics.finish!
+        configuration.notify(
+          error,
+          context: {
+            table_name: table_name,
+            job: job_class_name,
+            action: "run",
+            run_metadata: metrics.to_h
+          }
+        )
+        errors << error
+        raise unless continue_on_error
+
+        metrics
+      ensure
+        configuration.current_run_metrics = nil
+      end
+    end
+
+    private :plan_all, :emit_run_metadata, :run_single_table!
+  end
+end
+
+require_relative "partition_gardener/rails" if defined?(Rails::Railtie)

data/sig/partition_gardener.rbs

@@ -0,0 +1,19 @@
+module PartitionGardener
+  VERSION: String
+
+  MOVE_BATCH_SIZE: Integer
+  FUTURE_MONTH_PARTITION_ROW_THRESHOLD: Integer
+  DEFAULT_STATEMENT_TIMEOUT: ActiveSupport::Duration
+
+  def self.run!: (
+    ?statement_timeout: ActiveSupport::Duration,
+    ?job_class_name: String,
+    ?continue_on_error: bool,
+    ?dry_run: bool,
+    ?table_name: String?
+  ) -> untyped
+
+  def self.audit: (String table_name) -> untyped
+  def self.plan: (?table_name: String?, ?config: Hash[Symbol, untyped]?) -> untyped
+  def self.maintenance_for: (Hash[Symbol, untyped] config, ?job_class_name: String) -> untyped
+end