partition_gardener 0.3.0

data/lib/partition_gardener/plan_applier.rb

@@ -0,0 +1,289 @@
+module PartitionGardener
+  class PlanApplier
+    include Naming
+
+    def initialize(config, executor: nil, job_class_name: "PartitionGardener")
+      @config = config
+      @executor = executor || Executor.for_config(config)
+      @job_class_name = job_class_name
+    end
+
+    def apply!(plan)
+      DefaultPartition.ensure!(@config, executor: @executor)
+      attached = attached_tail_segments
+      return unless plan.changed?(attached)
+
+      table_name = @config[:table_name]
+      staging_name = rebalance_staging_partition_name(table_name)
+      plan_signature = PlanDiff.plan_signature(plan.segments)
+      PartitionGardener.configuration.current_run_metrics&.plan_signature = plan_signature
+      record = run_record_enabled? ? RunRecord.load(table_name) : nil
+
+      guard_against_orphaned_staging!(staging_name, record, plan_signature)
+      operations = PlanDiff.operations(attached, plan.segments)
+
+      if resumable_run?(record, plan_signature, staging_name)
+        notify_resume(record)
+        apply_resuming!(plan, operations, staging_name, record)
+      elsif incremental_rebalance?
+        apply_incremental!(plan, operations, staging_name, plan_signature)
+      else
+        apply_full!(plan, staging_name, plan_signature)
+      end
+
+      analyze_parent_table!(table_name)
+      notify_rebalance_complete(table_name, plan)
+      RunRecord.clear(table_name) if run_record_enabled?
+    end
+
+    private
+
+    def strategy
+      Strategy.for(@config)
+    end
+
+    def conflict_key
+      @config[:conflict_key] || begin
+        [@config[:partition_key_column].split("::").first.strip, "id"].uniq
+      end
+    end
+
+    def cursor_columns
+      strategy.cursor_columns
+    end
+
+    def connection
+      Connection.connection
+    end
+
+    def attached_tail_segments
+      Planner.new(@config).attached_tail_segments
+    end
+
+    def managed_tail_partition_names
+      strategy.managed_tail_partition_names
+    end
+
+    def incremental_rebalance?
+      @config.fetch(:incremental_rebalance, PartitionGardener.configuration.incremental_rebalance)
+    end
+
+    def run_record_enabled?
+      @config.fetch(:run_record_enabled, PartitionGardener.configuration.run_record_enabled)
+    end
+
+    def resumable_run?(record, plan_signature, staging_name)
+      return false unless run_record_enabled? && record&.incomplete?
+      return false unless record.plan_signature == plan_signature
+      return false unless Connection.partition_exists?(staging_name)
+
+      true
+    end
+
+    def guard_against_orphaned_staging!(staging_name, record, plan_signature)
+      return unless Connection.partition_exists?(staging_name)
+      return if resumable_run?(record, plan_signature, staging_name)
+
+      row_count = Connection.count_rows_in_partition_table(staging_name)
+      return if row_count.zero?
+
+      raise OrphanedRebalanceStaging,
+        "#{staging_name} holds #{row_count} row(s) from an interrupted rebalance; " \
+        "restore a matching run record or move rows manually before maintenance"
+    end
+
+    def apply_full!(plan, staging_name, plan_signature)
+      record = start_run_record(plan_signature) if run_record_enabled?
+
+      prepare_staging!(staging_name)
+      record&.advance!("detach")
+
+      detach_managed_tail_partitions!(staging_name)
+      drain_default_tail_rows_into_staging!(staging_name)
+      record&.advance!("detach", staging_row_count: staging_row_count(staging_name))
+
+      create_target_segments!(plan)
+      record&.advance!("segments")
+
+      move_staging_rows!(staging_name)
+      record&.advance!("rows", staging_row_count: 0)
+
+      @executor.drop_table(staging_name)
+      record&.advance!("cleanup")
+    end
+
+    def apply_incremental!(plan, operations, staging_name, plan_signature)
+      changed_operations = operations.reject { |operation| operation.action == :keep }
+      return if changed_operations.empty?
+
+      record = start_run_record(plan_signature) if run_record_enabled?
+      keep_names = operations.select { |operation| operation.action == :keep }.map { |operation| operation.segment.name }
+
+      prepare_staging!(staging_name)
+      record&.advance!("detach")
+
+      detach_managed_tail_partitions!(staging_name, skip_names: keep_names)
+      drain_default_tail_rows_into_staging!(staging_name)
+      record&.advance!("detach", staging_row_count: staging_row_count(staging_name))
+
+      segments_to_create = changed_operations.filter_map(&:segment).uniq(&:name)
+      create_segments!(segments_to_create)
+      record&.advance!("segments")
+
+      move_staging_rows!(staging_name)
+      record&.advance!("rows", staging_row_count: 0)
+
+      drop_removed_partitions!(operations)
+      @executor.drop_table(staging_name)
+      record&.advance!("cleanup")
+    end
+
+    def apply_resuming!(plan, operations, staging_name, record)
+      unless record.phase_at_least?("detach")
+        prepare_staging!(staging_name)
+        detach_managed_tail_partitions!(staging_name, skip_names: keep_names_from(operations))
+        drain_default_tail_rows_into_staging!(staging_name)
+        record = record.advance!("detach", staging_row_count: staging_row_count(staging_name))
+      end
+
+      unless record.phase_at_least?("segments")
+        create_target_segments!(plan)
+        record = record.advance!("segments")
+      end
+
+      unless record.phase_at_least?("rows")
+        move_staging_rows!(staging_name)
+        record = record.advance!("rows", staging_row_count: 0)
+      end
+
+      drop_removed_partitions!(operations) if incremental_rebalance?
+
+      unless record.phase_at_least?("cleanup")
+        @executor.drop_table(staging_name) if Connection.partition_exists?(staging_name)
+        record.advance!("cleanup")
+      end
+    end
+
+    def keep_names_from(operations)
+      operations.select { |operation| operation.action == :keep }.map { |operation| operation.segment.name }
+    end
+
+    def prepare_staging!(staging_name)
+      @executor.drop_table(staging_name)
+      @executor.ensure_detached_partition_table!(@config[:table_name], staging_name, conflict_key: conflict_key)
+    end
+
+    def create_target_segments!(plan)
+      create_segments!(plan.segments)
+    end
+
+    def create_segments!(segments)
+      return if MaintenanceBackend.hybrid?(@config)
+
+      table_name = @config[:table_name]
+
+      segments.each do |segment|
+        next if Connection.partition_attached?(table_name, segment.name)
+
+        @executor.create_partition(
+          table_name,
+          segment.name,
+          segment.for_values_clause(strategy)
+        )
+      end
+    end
+
+    def move_staging_rows!(staging_name)
+      table_name = @config[:table_name]
+      return unless Connection.partition_exists?(staging_name)
+      return if Connection.count_rows_in_partition_table(staging_name).zero?
+
+      @executor.move_all_rows_to_parent!(table_name, staging_name, conflict_key, cursor_columns: cursor_columns)
+    end
+
+    def drop_removed_partitions!(operations)
+      operations.select { |operation| operation.action == :drop }.each do |operation|
+        partition_name = operation.attached_segment.name
+        next if Connection.partition_attached?(@config[:table_name], partition_name)
+
+        @executor.drop_table(partition_name) if Connection.partition_exists?(partition_name)
+      end
+    end
+
+    def detach_managed_tail_partitions!(staging_name, skip_names: [])
+      table_name = @config[:table_name]
+
+      managed_tail_partition_names.each do |partition_name|
+        next if skip_names.include?(partition_name)
+        next unless Connection.partition_attached?(table_name, partition_name)
+
+        @executor.detach_partition(table_name, partition_name)
+        @executor.move_all_rows_between_partitions!(partition_name, staging_name, conflict_key, cursor_columns: cursor_columns)
+        @executor.drop_table(partition_name)
+      end
+    end
+
+    def drain_default_tail_rows_into_staging!(staging_name)
+      default_name = default_partition_name(@config[:table_name])
+      return unless Connection.partition_exists?(default_name)
+
+      where_condition = strategy.rebalance_default_drain_where_condition
+      return if where_condition == "FALSE"
+      return if Connection.count_rows_in_partition(default_name, where_condition).zero?
+
+      @executor.drain_rows_between_partitions!(
+        default_name,
+        staging_name,
+        where_condition,
+        conflict_key,
+        cursor_columns: cursor_columns
+      )
+    end
+
+    def staging_row_count(staging_name)
+      return 0 unless Connection.partition_exists?(staging_name)
+
+      Connection.count_rows_in_partition_table(staging_name)
+    end
+
+    def start_run_record(plan_signature)
+      RunRecord.start(table_name: @config[:table_name], plan_signature: plan_signature)
+    end
+
+    def notify_resume(record)
+      PartitionGardener.configuration.notify(
+        "[PartitionGardener] Resuming rebalance for #{@config[:table_name]} at phase #{record.phase}",
+        context: {
+          table_name: @config[:table_name],
+          job: @job_class_name,
+          action: "resume",
+          phase: record.phase,
+          plan_signature: record.plan_signature
+        }
+      )
+    end
+
+    def notify_rebalance_complete(table_name, plan)
+      PartitionGardener.configuration.notify(
+        "[PartitionGardener] Rebalanced tail partitions for #{table_name}",
+        context: {
+          table_name: table_name,
+          layout: @config.fetch(:layout, :sliding_window),
+          segments: plan.segments.map(&:signature),
+          hot_buckets: plan.hot_buckets,
+          incremental: incremental_rebalance?
+        }
+      )
+    end
+
+    def analyze_parent_table!(table_name)
+      return unless analyze_after_rebalance?
+
+      Connection.connection.execute("ANALYZE #{Connection.quoted_table(table_name)}")
+    end
+
+    def analyze_after_rebalance?
+      @config.fetch(:analyze_after_rebalance, PartitionGardener.configuration.analyze_after_rebalance)
+    end
+  end
+end

data/lib/partition_gardener/plan_diff.rb

@@ -0,0 +1,51 @@
+require "digest"
+require "json"
+
+module PartitionGardener
+  class PlanDiff
+    Operation = Data.define(:action, :segment, :attached_segment)
+
+    def self.operations(attached_segments, target_segments)
+      new(attached_segments, target_segments).operations
+    end
+
+    def self.plan_signature(segments)
+      payload = segments.map(&:signature).sort_by(&:to_s).to_json
+      Digest::SHA256.hexdigest(payload)[0, 16]
+    end
+
+    def self.changed?(attached_segments, target_segments)
+      attached_segments.map(&:signature).sort != target_segments.map(&:signature).sort
+    end
+
+    def initialize(attached_segments, target_segments)
+      @attached_segments = attached_segments
+      @target_segments = target_segments
+    end
+
+    def operations
+      attached_by_name = @attached_segments.each_with_object({}) { |segment, index| index[segment.name] = segment }
+      target_by_name = @target_segments.each_with_object({}) { |segment, index| index[segment.name] = segment }
+      result = []
+
+      target_by_name.each_value do |segment|
+        attached = attached_by_name[segment.name]
+        result << if attached.nil?
+          Operation.new(:create, segment, nil)
+        elsif attached.signature == segment.signature
+          Operation.new(:keep, segment, attached)
+        else
+          Operation.new(:reshape, segment, attached)
+        end
+      end
+
+      attached_by_name.each_value do |attached|
+        next if target_by_name[attached.name]
+
+        result << Operation.new(:drop, nil, attached)
+      end
+
+      result
+    end
+  end
+end

data/lib/partition_gardener/plan_report.rb

@@ -0,0 +1,95 @@
+require "json"
+
+module PartitionGardener
+  class PlanReport
+    SCHEMA_VERSION = "1.0"
+
+    def self.build(config)
+      planner = Planner.new(config)
+      plan = planner.build
+      attached_segments = planner.attached_tail_segments
+      operations = PlanDiff.operations(attached_segments, plan.segments)
+      gaps = GapDetection.call(config[:table_name], config: config)
+
+      new(
+        table_name: config[:table_name],
+        layout: config.fetch(:layout, :sliding_window),
+        changed: PlanDiff.changed?(attached_segments, plan.segments),
+        plan_signature: PlanDiff.plan_signature(plan.segments),
+        target_segments: plan.segments,
+        attached_segments: attached_segments,
+        operations: operations,
+        gaps: gaps,
+        hot_buckets: plan.hot_buckets
+      )
+    end
+
+    attr_reader :table_name,
+      :layout,
+      :changed,
+      :plan_signature,
+      :target_segments,
+      :attached_segments,
+      :operations,
+      :gaps,
+      :hot_buckets
+
+    def initialize(table_name:, layout:, changed:, plan_signature:, target_segments:, attached_segments:, operations:, gaps:, hot_buckets:)
+      @table_name = table_name
+      @layout = layout
+      @changed = changed
+      @plan_signature = plan_signature
+      @target_segments = target_segments
+      @attached_segments = attached_segments
+      @operations = operations
+      @gaps = gaps
+      @hot_buckets = hot_buckets
+    end
+
+    def to_h
+      {
+        schema_version: self.class::SCHEMA_VERSION,
+        table_name: table_name,
+        layout: layout,
+        changed: changed,
+        plan_signature: plan_signature,
+        target_segments: target_segments.map { |segment| segment_to_h(segment) },
+        attached_segments: attached_segments.map { |segment| segment_to_h(segment) },
+        operations: operations.map { |operation| operation_to_h(operation) },
+        gaps: gaps.map { |gap| gap_to_h(gap) },
+        hot_buckets: hot_buckets
+      }
+    end
+
+    def to_json(...)
+      to_h.to_json(...)
+    end
+
+    private
+
+    def segment_to_h(segment)
+      {
+        name: segment.name,
+        range_start: segment.range_start,
+        range_end: segment.range_end,
+        kind: segment.kind
+      }
+    end
+
+    def operation_to_h(operation)
+      {
+        action: operation.action,
+        segment: operation.segment ? segment_to_h(operation.segment) : nil,
+        attached_segment: operation.attached_segment ? segment_to_h(operation.attached_segment) : nil
+      }
+    end
+
+    def gap_to_h(gap)
+      {
+        range_start: gap.range_start,
+        range_end: gap.range_end,
+        message: gap.message
+      }
+    end
+  end
+end

data/lib/partition_gardener/planner.rb

@@ -0,0 +1,21 @@
+module PartitionGardener
+  class Planner
+    def initialize(config)
+      @config = config
+    end
+
+    def build
+      strategy.build_plan
+    end
+
+    def attached_tail_segments
+      strategy.attached_tail_segments
+    end
+
+    private
+
+    def strategy
+      Strategy.for(@config)
+    end
+  end
+end

data/lib/partition_gardener/predicate.rb

@@ -0,0 +1,85 @@
+module PartitionGardener
+  module Predicate
+    COLUMN_PATTERN = /\A[a-z_][a-z0-9_]*\z/i
+    OPERATORS = %w[eq ne is_null is_not_null].freeze
+
+    module_function
+
+    def render(predicate, connection: nil)
+      connection ||= Connection.connection
+      predicate = predicate.transform_keys(&:to_sym)
+      column = predicate.fetch(:column)
+      operator = predicate.fetch(:operator).to_s
+
+      validate_column!(column)
+      validate_operator!(operator)
+
+      quoted_column = connection.quote_column_name(column)
+
+      case operator
+      when "eq"
+        "#{quoted_column} = #{connection.quote(predicate.fetch(:value))}"
+      when "ne"
+        "#{quoted_column} <> #{connection.quote(predicate.fetch(:value))}"
+      when "is_null"
+        "#{quoted_column} IS NULL"
+      when "is_not_null"
+        "#{quoted_column} IS NOT NULL"
+      end
+    end
+
+    def normalize_branch!(branch, discriminator_column: nil, connection: nil)
+      entry = branch.transform_keys(&:to_sym)
+      entry = entry.dup
+
+      if entry[:predicate]
+        entry[:where_condition] = render(entry[:predicate], connection: connection)
+        entry.delete(:predicate)
+      elsif entry[:where_condition]
+        entry[:where_condition] = entry[:where_condition].to_s
+      elsif (column = entry[:discriminator_column] || discriminator_column) && entry.key?(:value)
+        entry[:where_condition] = render(
+          {column: column, operator: "eq", value: entry[:value]},
+          connection: connection
+        )
+      else
+        name = entry[:name] || "?"
+        raise ArgumentError,
+          "branch #{name.inspect} needs predicate, where_condition, or value with discriminator_column"
+      end
+
+      entry
+    end
+
+    def normalize_branches!(branches, discriminator_column: nil, connection: nil)
+      branches.map do |branch|
+        normalize_branch!(branch, discriminator_column: discriminator_column, connection: connection)
+      end
+    end
+
+    def list_branch_entries(branches, discriminator_column: nil, connection: nil)
+      normalize_branches!(branches, discriminator_column: discriminator_column, connection: connection).map do |branch|
+        {
+          name: branch.fetch(:name),
+          value: branch.fetch(:value),
+          where_condition: branch.fetch(:where_condition)
+        }
+      end
+    end
+
+    def validate_column!(column)
+      name = column.to_s
+      return if COLUMN_PATTERN.match?(name)
+
+      raise ArgumentError, "predicate column must be a simple identifier, got #{name.inspect}"
+    end
+
+    def validate_operator!(operator)
+      return if OPERATORS.include?(operator)
+
+      raise ArgumentError, "unsupported predicate operator: #{operator.inspect} (allowed: #{OPERATORS.join(", ")})"
+    end
+
+    private :validate_column!, :validate_operator!
+  end
+end

data/lib/partition_gardener/premake_monthly_maintenance.rb

@@ -0,0 +1,44 @@
+module PartitionGardener
+  class PremakeMonthlyMaintenance < DateRangeMaintenance
+    def run!
+      report_audit_warnings
+      ensure_default_partition
+      ensure_premade_months!
+      apply_archive_retention!
+      drain_default_partition
+    end
+
+    private
+
+    def ensure_premade_months!
+      return unless Connection.table_is_partitioned?(table_name)
+
+      month_count = @config.fetch(:premake_months, 3)
+      start_month = DateCalendar.beginning_of_month(PartitionGardener.configuration.today)
+
+      (0..month_count).each do |offset|
+        identifier = DateCalendar.add_months(start_month, offset)
+        ensure_month_partition!(identifier)
+      end
+    end
+
+    def ensure_month_partition!(identifier)
+      partition_name = @config[:partition_name_format].call(identifier)
+      for_values_clause = @config[:partition_definition].call(identifier)
+
+      return if Connection.partition_attached?(table_name, partition_name)
+
+      if Connection.partition_exists?(partition_name)
+        send(
+          :attach_archive_partition!,
+          table_name,
+          partition_name,
+          for_values_clause,
+          strategy.bucket_where_condition(identifier)
+        )
+      else
+        @executor.create_partition(table_name, partition_name, for_values_clause)
+      end
+    end
+  end
+end

data/lib/partition_gardener/rails.rb

@@ -0,0 +1,12 @@
+require "rails"
+
+module PartitionGardener
+  class Railtie < Rails::Railtie
+    initializer "partition_gardener.connection" do
+      PartitionGardener.configure do |config|
+        config.connection_resolver = -> { ActiveRecord::Base.connection }
+        config.today_resolver = -> { Time.zone.today }
+      end
+    end
+  end
+end

data/lib/partition_gardener/registry.rb

@@ -0,0 +1,84 @@
+module PartitionGardener
+  module Registry
+    class << self
+      def register(config)
+        normalized = Templates.normalize(config)
+        normalized[:maintenance_backend] = MaintenanceBackend.normalize(normalized[:maintenance_backend])
+        tables.delete_if { |entry| entry[:table_name] == normalized[:table_name] }
+        tables << normalized
+        MaintenanceBackend.validate!(normalized)
+        normalized
+      end
+
+      def register_template(template_name, **options)
+        builder = Templates.public_method(template_name)
+        register(builder.call(**options))
+      rescue NameError
+        raise ArgumentError, "unknown template: #{template_name.inspect}"
+      end
+
+      def register_all(configs)
+        configs.map { |config| register(config) }
+      end
+
+      def each_table_config(&block)
+        expanded_table_configs.each(&block)
+      end
+
+      def expanded_table_configs
+        expanded_tables
+      end
+
+      def configs_for_table(table_name)
+        name = table_name.to_s
+        registered = tables.find { |config| config[:table_name].to_s == name }
+        return Templates.expand(registered) if registered&.fetch(:layout, nil) == :composite
+
+        expanded_table_configs.select do |config|
+          config[:table_name].to_s == name || config[:parent_table_name].to_s == name
+        end
+      end
+
+      def find_by_table_name(table_name)
+        expanded_tables.find { |config| config[:table_name].to_s == table_name.to_s }
+      end
+
+      def hot_switch_partition_config(table_name)
+        config = find_by_table_name(table_name)
+        return nil unless config
+
+        {
+          table_name: config[:table_name],
+          partition_name_format: config[:partition_name_format],
+          partition_definition: config[:partition_definition],
+          partitions_to_create: lambda { |today|
+            [DateCalendar.beginning_of_month(today), DateCalendar.beginning_of_month(DateCalendar.next_month(today))]
+          }
+        }
+      end
+
+      def tables
+        @tables ||= []
+      end
+
+      def map(&block)
+        each_table_config.map(&block)
+      end
+
+      # Prefer `.tables`. Kept for callers that used the old constant name.
+      def self.TABLES
+        tables
+      end
+
+      def reset!
+        @tables = []
+      end
+
+      private
+
+      def expanded_tables
+        tables.flat_map { |config| Templates.expand(config) }
+      end
+    end
+  end
+end

data/lib/partition_gardener/run_failed.rb

@@ -0,0 +1,10 @@
+module PartitionGardener
+  class RunFailed < StandardError
+    attr_reader :errors
+
+    def initialize(errors)
+      @errors = errors
+      super("Partition maintenance failed for #{errors.size} table(s)")
+    end
+  end
+end