partition_gardener 0.3.0

data/lib/partition_gardener/date_range_maintenance.rb

@@ -0,0 +1,297 @@
+module PartitionGardener
+  class DateRangeMaintenance
+    include Naming
+
+    def initialize(config, job_class_name: "PartitionGardener", executor: nil)
+      @config = config
+      @job_class_name = job_class_name
+      @executor = executor || Executor.for_config(config)
+    end
+
+    def run!
+      report_audit_warnings
+      ensure_default_partition
+      unless MaintenanceBackend.hybrid?(@config)
+        finalize_archive_partitions
+        apply_archive_retention!
+      end
+      rebalance_tail!
+      drain_default_partition
+    end
+
+    def split_future_month_from_current!(_identifier = nil)
+      rebalance_tail!
+    end
+
+    def split_pressured_future_month_partitions
+      rebalance_tail!
+    end
+
+    def collapse_low_volume_future_month_partitions
+      rebalance_tail!
+    end
+
+    private
+
+    def strategy
+      Strategy.for(@config)
+    end
+
+    def connection
+      Connection.connection
+    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 table_name
+      @config[:table_name]
+    end
+
+    def notify_error(error, action)
+      PartitionGardener.configuration.notify(
+        error,
+        context: {
+          table_name: table_name,
+          job: @job_class_name,
+          action: action
+        }
+      )
+    end
+
+    def rebalance_tail!
+      return unless Connection.table_is_partitioned?(table_name)
+
+      plan = strategy.build_plan
+      PlanApplier.new(@config, executor: @executor, job_class_name: @job_class_name).apply!(plan)
+    rescue => error
+      notify_error(error, "rebalance_tail")
+      raise
+    end
+
+    def ensure_default_partition
+      DefaultPartition.ensure!(@config, executor: @executor)
+    rescue => error
+      notify_error(error, "ensure_default_partition")
+      raise
+    end
+
+    def finalize_archive_partitions
+      return unless Connection.table_is_partitioned?(table_name)
+      return if strategy.is_a?(Strategy::HashBranches)
+      return if strategy.is_a?(Strategy::ListSplit)
+      return if strategy.is_a?(Strategy::IntegerRange)
+
+      source_name = archive_finalize_source_partition_name
+      return unless source_name
+      return unless Connection.partition_exists?(source_name)
+
+      Connection.get_distinct_partition_identifiers(
+        source_name,
+        @config[:partition_key_column],
+        @config[:extract_partition_identifier]
+      ).each do |identifier|
+        next unless strategy.archive_bucket?(identifier)
+
+        finalize_archive_from_source!(identifier, source_name)
+      end
+    rescue => error
+      notify_error(error, "finalize_archive_partitions")
+      raise
+    end
+
+    def archive_finalize_source_partition_name
+      current_name = current_partition_name(table_name)
+      return current_name if Connection.partition_exists?(current_name)
+
+      default_partition_name(table_name) if Connection.partition_exists?(default_partition_name(table_name))
+    end
+
+    def finalize_archive_from_source!(identifier, source_partition_name)
+      partition_name = archive_partition_name(identifier)
+      where_condition = strategy.bucket_where_condition(identifier)
+
+      ensure_archive_partition_attached!(identifier, partition_name: partition_name, source_partition_name: source_partition_name)
+
+      record_count = Connection.count_rows_in_partition(source_partition_name, where_condition)
+      return if record_count.zero?
+
+      @executor.move_rows_to_parent_partition!(
+        table_name: table_name,
+        source_partition_name: source_partition_name,
+        where_condition: where_condition,
+        destination_partition_name: partition_name,
+        conflict_key: conflict_key,
+        record_count: record_count,
+        cursor_columns: cursor_columns
+      )
+    end
+
+    def archive_partition_name(identifier)
+      if @config[:archive_partition_name_format]
+        @config[:archive_partition_name_format].call(identifier)
+      else
+        @config[:partition_name_format].call(identifier)
+      end
+    end
+
+    def drain_default_partition
+      return unless Connection.table_is_partitioned?(table_name)
+
+      default_name = default_partition_name(table_name)
+      return unless Connection.partition_exists?(default_name)
+
+      default_row_count = Connection.count_rows_in_partition_table(default_name)
+      return if default_row_count.zero?
+
+      PartitionGardener.configuration.notify(
+        "[PartitionGardener] Found #{default_row_count} rows in #{default_name}",
+        context: {
+          table_name: table_name,
+          default_partition_name: default_name,
+          count: default_row_count
+        }
+      )
+
+      unless strategy.is_a?(Strategy::HashBranches) || strategy.is_a?(Strategy::ListSplit)
+        bucket_counts(default_name).each do |identifier, row_count|
+          if strategy.archive_bucket?(identifier)
+            finalize_archive_from_source!(identifier, default_name)
+          elsif strategy.future_bucket?(identifier) && row_count >= split_row_threshold
+            finalize_archive_from_source!(identifier, default_name)
+          end
+        end
+      end
+
+      skip_bucket_drain = strategy.is_a?(Strategy::HashBranches) || strategy.is_a?(Strategy::ListSplit)
+      drain_remaining_default_rows!(known_row_count: skip_bucket_drain ? default_row_count : nil)
+    rescue => error
+      notify_error(error, "drain_default_partition")
+      raise
+    end
+
+    def drain_remaining_default_rows!(known_row_count: nil)
+      default_name = default_partition_name(table_name)
+      where_condition = strategy.default_partition_drain_where_condition
+      return if where_condition == "FALSE"
+
+      record_count = known_row_count || Connection.count_rows_in_partition(default_name, where_condition)
+      return if record_count.zero?
+
+      @executor.move_rows_to_parent_partition!(
+        table_name: table_name,
+        source_partition_name: default_name,
+        where_condition: where_condition,
+        destination_partition_name: table_name,
+        conflict_key: conflict_key,
+        record_count: record_count,
+        cursor_columns: cursor_columns
+      )
+    end
+
+    def split_row_threshold
+      @config.fetch(:split_row_threshold, FUTURE_MONTH_PARTITION_ROW_THRESHOLD)
+    end
+
+    def bucket_counts(partition_name)
+      if strategy.is_a?(Strategy::DateRange)
+        strategy.bucket_counts_in_partition(partition_name)
+      else
+        Connection.get_distinct_partition_identifiers(
+          partition_name,
+          @config[:partition_key_column],
+          @config[:extract_partition_identifier]
+        ).index_with do |identifier|
+          Connection.count_rows_in_partition(partition_name, strategy.bucket_where_condition(identifier))
+        end
+      end
+    end
+
+    def ensure_archive_partition_attached!(identifier, partition_name:, source_partition_name:)
+      for_values_clause = archive_for_values_clause(identifier)
+      where_condition = strategy.bucket_where_condition(identifier)
+
+      return if Connection.partition_attached?(table_name, partition_name)
+
+      if Connection.count_rows_in_partition(source_partition_name, where_condition).positive?
+        @executor.ensure_detached_partition_table!(table_name, partition_name, conflict_key: conflict_key)
+        @executor.drain_rows_between_partitions!(
+          source_partition_name,
+          partition_name,
+          where_condition,
+          conflict_key,
+          cursor_columns: cursor_columns
+        )
+      end
+
+      if Connection.partition_exists?(partition_name)
+        attach_archive_partition!(table_name, partition_name, for_values_clause, where_condition)
+      else
+        @executor.create_partition(table_name, partition_name, for_values_clause)
+      end
+    end
+
+    def attach_archive_partition!(table_name, partition_name, for_values_clause, where_condition)
+      @executor.attach_partition(table_name, partition_name, for_values_clause)
+    rescue => error
+      raise unless check_violation?(error)
+
+      default_name = default_partition_name(table_name)
+      if Connection.partition_exists?(default_name) &&
+          Connection.count_rows_in_partition(default_name, where_condition).positive?
+        @executor.drain_rows_between_partitions!(
+          default_name,
+          partition_name,
+          where_condition,
+          conflict_key,
+          cursor_columns: cursor_columns
+        )
+      end
+
+      @executor.attach_partition(table_name, partition_name, for_values_clause)
+    end
+
+    def check_violation?(error)
+      message = error.message.to_s
+      return true if message.include?("CheckViolation") || message.include?("check constraint")
+
+      error.cause && check_violation?(error.cause)
+    end
+
+    def report_audit_warnings
+      audit = Audit.call(table_name, config: @config)
+      audit.warnings.each do |warning|
+        PartitionGardener.configuration.notify(
+          "[PartitionGardener] #{warning}",
+          context: {
+            table_name: table_name,
+            job: @job_class_name,
+            action: "audit"
+          }
+        )
+      end
+    end
+
+    def apply_archive_retention!
+      ArchiveRetention.new(@config, executor: @executor, job_class_name: @job_class_name).apply!
+    end
+
+    def archive_for_values_clause(identifier)
+      if strategy.is_a?(Strategy::IntegerRange)
+        band = @config.fetch(:archive_band_size, Strategy::IntegerRange::DEFAULT_ARCHIVE_BAND_SIZE)
+        "FROM (#{identifier}) TO (#{identifier + band})"
+      else
+        @config[:partition_definition].call(identifier)
+      end
+    end
+  end
+
+  ThreeAreaMaintenance = DateRangeMaintenance
+end

data/lib/partition_gardener/default_partition.rb

@@ -0,0 +1,54 @@
+module PartitionGardener
+  class MissingDefaultPartition < StandardError; end
+
+  class DefaultPartition
+    include Naming
+
+    def self.ensure!(config, executor: nil)
+      new(config, executor: executor).ensure!
+    end
+
+    def initialize(config, executor: nil)
+      @config = config
+      @executor = executor || Executor.for_config(config)
+    end
+
+    def ensure!
+      table_name = @config[:table_name]
+      return unless Connection.table_is_partitioned?(table_name)
+      return unless Strategy.for(@config).default_partition_required?
+
+      default_name = default_partition_name(table_name)
+
+      if Connection.partition_exists?(default_name)
+        attach_existing_default!(table_name, default_name)
+      else
+        create_default!(table_name, default_name)
+      end
+
+      return if Connection.partition_attached?(table_name, default_name)
+
+      raise MissingDefaultPartition,
+        "Default partition #{default_name} is required for #{table_name} but is not attached"
+    end
+
+    private
+
+    def connection
+      Connection.connection
+    end
+
+    def attach_existing_default!(table_name, default_name)
+      return if Connection.partition_attached?(table_name, default_name)
+
+      @executor.attach_default_partition(table_name, default_name)
+    end
+
+    def create_default!(table_name, default_name)
+      sql = <<~SQL
+        CREATE TABLE IF NOT EXISTS #{Connection.quoted_table(default_name)} PARTITION OF #{Connection.quoted_table(table_name)} DEFAULT
+      SQL
+      connection.execute(sql)
+    end
+  end
+end

data/lib/partition_gardener/executor.rb

@@ -0,0 +1,324 @@
+require "json"
+
+module PartitionGardener
+  class Executor
+    def self.for_config(config, connection: Connection.connection)
+      new(
+        connection: connection,
+        batch_size: config.fetch(:move_batch_size, MOVE_BATCH_SIZE)
+      )
+    end
+
+    def initialize(connection: Connection.connection, batch_size: MOVE_BATCH_SIZE)
+      @connection = connection
+      @batch_size = batch_size
+    end
+
+    def attach_partition(table_name, partition_name, for_values_clause)
+      sql = <<~SQL
+        ALTER TABLE #{quoted_table(table_name)} ATTACH PARTITION #{quoted_table(partition_name)}
+        FOR VALUES #{for_values_clause}
+      SQL
+      @connection.execute(sql)
+    end
+
+    def attach_default_partition(table_name, partition_name)
+      sql = <<~SQL
+        ALTER TABLE #{quoted_table(table_name)} ATTACH PARTITION #{quoted_table(partition_name)} DEFAULT
+      SQL
+      @connection.execute(sql)
+    end
+
+    def detach_partition(table_name, partition_name, concurrently: false)
+      concurrently_keyword = concurrently ? "CONCURRENTLY " : ""
+      sql = <<~SQL
+        ALTER TABLE #{quoted_table(table_name)} DETACH PARTITION #{concurrently_keyword}#{quoted_table(partition_name)}
+      SQL
+      @connection.execute(sql)
+    end
+
+    def drop_table(table_name)
+      sql = <<~SQL
+        DROP TABLE IF EXISTS #{quoted_table(table_name)}
+      SQL
+      @connection.execute(sql)
+    end
+
+    def create_partition(table_name, partition_name, for_values_clause)
+      sql = <<~SQL
+        CREATE TABLE IF NOT EXISTS #{quoted_table(partition_name)} PARTITION OF #{quoted_table(table_name)}
+        FOR VALUES #{for_values_clause}
+      SQL
+      @connection.execute(sql)
+    end
+
+    def ensure_detached_partition_table!(table_name, partition_name, conflict_key:)
+      sql = <<~SQL
+        CREATE TABLE IF NOT EXISTS #{quoted_table(partition_name)} (
+          LIKE #{quoted_table(table_name)} INCLUDING ALL
+        )
+      SQL
+      @connection.execute(sql)
+
+      conflict_columns = conflict_key.map { |column| @connection.quote_column_name(column) }.join(", ")
+      sql = <<~SQL
+        CREATE UNIQUE INDEX IF NOT EXISTS #{@connection.quote_column_name("#{partition_name}_conflict_key_idx")}
+        ON #{quoted_table(partition_name)} (#{conflict_columns})
+      SQL
+      @connection.execute(sql)
+    end
+
+    def move_all_rows_between_partitions!(source_partition_name, destination_partition_name, conflict_key, cursor_columns: conflict_key)
+      move_rows_between_partitions!(
+        source_partition_name: source_partition_name,
+        destination_partition_name: destination_partition_name,
+        where_condition: nil,
+        conflict_key: conflict_key,
+        cursor_columns: cursor_columns
+      )
+    end
+
+    def drain_rows_between_partitions!(source_partition_name, destination_partition_name, where_condition, conflict_key, cursor_columns: conflict_key)
+      move_rows_between_partitions!(
+        source_partition_name: source_partition_name,
+        destination_partition_name: destination_partition_name,
+        where_condition: where_condition,
+        conflict_key: conflict_key,
+        cursor_columns: cursor_columns
+      )
+    end
+
+    def move_all_rows_to_parent!(table_name, source_partition_name, conflict_key, cursor_columns: conflict_key)
+      ensure_parent_conflict_index!(table_name, conflict_key)
+      move_rows_with_keyset!(
+        source_partition_name: source_partition_name,
+        insert_target: quoted_table(table_name),
+        where_condition: nil,
+        conflict_key: conflict_key,
+        cursor_columns: cursor_columns
+      )
+    end
+
+    def move_rows_to_parent_partition!(
+      table_name:,
+      source_partition_name:,
+      where_condition:,
+      destination_partition_name:,
+      conflict_key:,
+      record_count:,
+      cursor_columns: conflict_key
+    )
+      ensure_parent_conflict_index!(table_name, conflict_key)
+
+      PartitionGardener.configuration.notify(
+        "[PartitionGardener] Moving #{record_count} rows from #{source_partition_name} to #{destination_partition_name}",
+        context: {
+          table_name: table_name,
+          source_partition_name: source_partition_name,
+          destination_partition_name: destination_partition_name,
+          record_count: record_count
+        }
+      )
+
+      moved_rows = move_rows_with_keyset!(
+        source_partition_name: source_partition_name,
+        insert_target: quoted_table(table_name),
+        where_condition: where_condition,
+        conflict_key: conflict_key,
+        cursor_columns: cursor_columns
+      )
+
+      PartitionGardener.configuration.notify(
+        "[PartitionGardener] Moved #{moved_rows} rows from #{source_partition_name} to #{destination_partition_name}",
+        context: {
+          table_name: table_name,
+          source_partition_name: source_partition_name,
+          destination_partition_name: destination_partition_name,
+          moved_rows: moved_rows
+        }
+      )
+    end
+
+    private
+
+    def quoted_table(name)
+      @connection.quote_table_name(name)
+    end
+
+    def quoted_conflict_key(conflict_key)
+      conflict_key.map { |column| @connection.quote_column_name(column) }.join(", ")
+    end
+
+    def move_rows_between_partitions!(source_partition_name:, destination_partition_name:, where_condition:, conflict_key:, cursor_columns:)
+      move_rows_with_keyset!(
+        source_partition_name: source_partition_name,
+        insert_target: quoted_table(destination_partition_name),
+        where_condition: where_condition,
+        conflict_key: conflict_key,
+        cursor_columns: cursor_columns
+      )
+    end
+
+    def move_rows_with_keyset!(source_partition_name:, insert_target:, where_condition:, conflict_key:, cursor_columns:)
+      moved_rows = 0
+      last_cursor = nil
+
+      loop do
+        batch_result = execute_move_batch(
+          source_partition_name: source_partition_name,
+          insert_target: insert_target,
+          where_condition: where_condition,
+          conflict_key: conflict_key,
+          cursor_columns: cursor_columns,
+          last_cursor: last_cursor
+        )
+
+        moved_rows += batch_result[:deleted]
+        last_cursor = batch_result[:last_cursor]
+        if batch_result[:deleted].zero?
+          if batch_result[:batch_size].positive?
+            raise UnmovedRowsRemaining.new(
+              source_partition_name: source_partition_name,
+              batch_size: batch_result[:batch_size],
+              last_cursor: last_cursor
+            )
+          end
+
+          break
+        end
+        break if batch_result[:batch_size] < @batch_size
+      end
+
+      record_rows_moved!(moved_rows)
+      moved_rows
+    end
+
+    def record_rows_moved!(count)
+      PartitionGardener.configuration.current_run_metrics&.add_rows(count)
+    end
+
+    def execute_move_batch(source_partition_name:, insert_target:, where_condition:, conflict_key:, cursor_columns:, last_cursor:)
+      order_sql = cursor_columns.map { |column| @connection.quote_column_name(column) }.join(", ")
+      conflict_key_sql = quoted_conflict_key(conflict_key)
+      cursor_columns_sql = order_sql
+      inserted_not_exists_sql = column_match_sql("inserted_rows", "batch_rows", cursor_columns)
+      target_exists_sql = column_match_sql("target_rows", "batch_rows", conflict_key)
+      removable_match_sql = column_match_sql("source_rows", "removable_rows", cursor_columns)
+      returning_sql = cursor_columns.map { |column| "source_rows.#{@connection.quote_column_name(column)}" }.join(", ")
+
+      where_parts = []
+      where_parts << where_condition if where_condition
+      keyset_sql = keyset_after_clause(cursor_columns, last_cursor)
+      where_parts << keyset_sql if keyset_sql
+      where_clause = where_parts.any? ? "WHERE #{where_parts.join(" AND ")}" : ""
+
+      sql = <<~SQL
+        WITH batch_rows AS (
+          SELECT *
+          FROM #{quoted_table(source_partition_name)}
+          #{where_clause}
+          ORDER BY #{order_sql}
+          LIMIT #{@batch_size}
+        ),
+        inserted_rows AS (
+          INSERT INTO #{insert_target}
+          SELECT * FROM batch_rows
+          ON CONFLICT (#{conflict_key_sql}) DO NOTHING
+          RETURNING #{cursor_columns_sql}
+        ),
+        duplicates_at_target AS (
+          SELECT #{cursor_columns_sql}
+          FROM batch_rows
+          WHERE NOT EXISTS (
+            SELECT 1
+            FROM inserted_rows
+            WHERE #{inserted_not_exists_sql}
+          )
+          AND EXISTS (
+            SELECT 1
+            FROM #{insert_target} AS target_rows
+            WHERE #{target_exists_sql}
+          )
+        ),
+        removable_rows AS (
+          SELECT #{cursor_columns_sql} FROM inserted_rows
+          UNION
+          SELECT #{cursor_columns_sql} FROM duplicates_at_target
+        ),
+        deleted_rows AS (
+          DELETE FROM #{quoted_table(source_partition_name)} AS source_rows
+          USING removable_rows
+          WHERE #{removable_match_sql}
+          RETURNING #{returning_sql}
+        )
+        SELECT
+          (SELECT COUNT(*)::int FROM deleted_rows) AS deleted,
+          (SELECT COUNT(*)::int FROM batch_rows) AS batch_size,
+          (
+            SELECT row_to_json(last_row)
+            FROM (
+              SELECT #{order_sql}
+              FROM batch_rows
+              ORDER BY #{order_sql} DESC
+              LIMIT 1
+            ) AS last_row
+          ) AS last_cursor
+      SQL
+
+      row = @connection.execute(sql).first
+      deleted = row["deleted"].to_i
+      batch_size = row["batch_size"].to_i
+      last_cursor_values = parse_last_cursor(row["last_cursor"], cursor_columns)
+
+      {
+        deleted: deleted,
+        batch_size: batch_size,
+        last_cursor: last_cursor_values
+      }
+    end
+
+    def column_match_sql(left_alias, right_alias, columns)
+      columns.map do |column|
+        quoted_column = @connection.quote_column_name(column)
+        "#{left_alias}.#{quoted_column} = #{right_alias}.#{quoted_column}"
+      end.join(" AND ")
+    end
+
+    def keyset_after_clause(columns, last_values)
+      return nil if last_values.nil?
+
+      column_sql = columns.map { |column| @connection.quote_column_name(column) }.join(", ")
+      value_sql = last_values.map { |value| @connection.quote(value) }.join(", ")
+      "(#{column_sql}) > (#{value_sql})"
+    end
+
+    def parse_last_cursor(last_cursor_json, cursor_columns)
+      return nil if last_cursor_json.nil?
+
+      payload = last_cursor_json.is_a?(String) ? JSON.parse(last_cursor_json) : last_cursor_json
+      cursor_columns.map { |column| payload[column] || payload[column.to_s] }
+    end
+
+    def ensure_parent_conflict_index!(table_name, conflict_key)
+      return if Connection.unique_index_covers?(table_name, conflict_key)
+
+      index_name = "#{table_name}_#{conflict_key.join("_")}_maintenance_conflict_idx"
+      columns_sql = conflict_key.map { |column| @connection.quote_column_name(column) }.join(", ")
+      sql = <<~SQL
+        CREATE UNIQUE INDEX IF NOT EXISTS #{@connection.quote_column_name(index_name)}
+        ON #{quoted_table(table_name)} (#{columns_sql})
+      SQL
+      @connection.execute(sql)
+
+      return if Connection.unique_index_covers?(table_name, conflict_key)
+
+      raise MissingConflictIndex,
+        "#{table_name} needs a unique index on (#{conflict_key.join(", ")}) including the partition key for row moves"
+    rescue => error
+      raise if error.is_a?(MissingConflictIndex)
+
+      raise MissingConflictIndex,
+        "#{table_name} needs a unique index on (#{conflict_key.join(", ")}) including the partition key for row moves: #{error.message}"
+    end
+  end
+end