partition_gardener 0.3.0

data/lib/partition_gardener/run_metrics.rb

@@ -0,0 +1,65 @@
+module PartitionGardener
+  class RunMetrics
+    attr_reader :table_name, :started_at, :rows_moved, :finished_at, :skipped, :skip_reason
+    attr_accessor :plan_signature
+
+    def initialize(table_name)
+      @table_name = table_name
+      @started_at = monotonic_clock
+      @plan_signature = nil
+      @rows_moved = 0
+      @finished_at = nil
+      @skipped = false
+      @skip_reason = nil
+    end
+
+    def add_rows(count)
+      @rows_moved += count.to_i
+    end
+
+    def mark_skipped!(reason)
+      @skipped = true
+      @skip_reason = reason
+      finish!
+    end
+
+    def finish!
+      @finished_at = monotonic_clock
+    end
+
+    def duration_ms
+      return nil unless @finished_at
+
+      ((@finished_at - @started_at) * 1000).round
+    end
+
+    def to_h
+      {
+        table_name: table_name,
+        duration_ms: duration_ms,
+        plan_signature: plan_signature,
+        rows_moved: rows_moved,
+        skipped: skipped,
+        skip_reason: skip_reason
+      }
+    end
+
+    private
+
+    def monotonic_clock
+      Process.clock_gettime(Process::CLOCK_MONOTONIC)
+    end
+  end
+
+  RunSummary = Data.define(:tables, :errors) do
+    SCHEMA_VERSION = "1.0"
+
+    def to_h
+      {
+        schema_version: SCHEMA_VERSION,
+        tables: tables.map(&:to_h),
+        errors: errors.map { |error| error.message }
+      }
+    end
+  end
+end

data/lib/partition_gardener/run_record.rb

@@ -0,0 +1,76 @@
+module PartitionGardener
+  class RunRecord
+    PHASES = %w[detach segments rows cleanup complete].freeze
+
+    attr_reader :table_name, :phase, :plan_signature, :staging_row_count
+
+    def self.load(table_name)
+      attributes = PartitionGardener.configuration.run_record_store.load(table_name)
+      return unless attributes
+
+      from_h(attributes)
+    end
+
+    def self.from_h(attributes)
+      new(
+        table_name: fetch_attribute(attributes, :table_name),
+        phase: fetch_attribute(attributes, :phase),
+        plan_signature: fetch_attribute(attributes, :plan_signature),
+        staging_row_count: fetch_attribute(attributes, :staging_row_count) || 0
+      )
+    end
+
+    def self.fetch_attribute(attributes, key)
+      attributes[key] || attributes[key.to_s]
+    end
+
+    def self.clear(table_name)
+      PartitionGardener.configuration.run_record_store.clear(table_name)
+    end
+
+    def self.start(table_name:, plan_signature:)
+      new(
+        table_name: table_name,
+        phase: "detach",
+        plan_signature: plan_signature,
+        staging_row_count: 0
+      ).tap(&:save!)
+    end
+
+    def initialize(table_name:, phase:, plan_signature:, staging_row_count: 0)
+      @table_name = table_name
+      @phase = phase
+      @plan_signature = plan_signature
+      @staging_row_count = staging_row_count
+    end
+
+    def incomplete?
+      Blank.present?(phase) && phase != "complete"
+    end
+
+    def save!
+      PartitionGardener.configuration.run_record_store.save(
+        table_name,
+        {
+          table_name: table_name,
+          phase: phase,
+          plan_signature: plan_signature,
+          staging_row_count: staging_row_count
+        }
+      )
+    end
+
+    def advance!(phase, staging_row_count: nil)
+      RunRecord.new(
+        table_name: table_name,
+        phase: phase,
+        plan_signature: plan_signature,
+        staging_row_count: staging_row_count.nil? ? self.staging_row_count : staging_row_count
+      ).tap(&:save!)
+    end
+
+    def phase_at_least?(name)
+      PHASES.index(phase) >= PHASES.index(name)
+    end
+  end
+end

data/lib/partition_gardener/sql_run_record_store.rb

@@ -0,0 +1,106 @@
+module PartitionGardener
+  class SqlRunRecordStore
+    TABLE_NAME = "partition_gardener_run_records"
+
+    def initialize
+      @schema_mutex = Mutex.new
+      @schema_ready = false
+    end
+
+    def load(table_name)
+      ensure_schema!
+      sql = load_sql(table_name)
+      row = connection.execute(sql).first
+
+      return unless row
+
+      {
+        table_name: row["table_name"],
+        phase: row["phase"],
+        plan_signature: row["plan_signature"],
+        staging_row_count: row["staging_row_count"].to_i
+      }
+    end
+
+    def save(table_name, attributes)
+      ensure_schema!
+      sql = save_sql(table_name, attributes)
+      connection.execute(sql)
+    end
+
+    def clear(table_name)
+      ensure_schema!
+      sql = clear_sql(table_name)
+      connection.execute(sql)
+    end
+
+    def ensure_schema!
+      return if @schema_ready
+
+      @schema_mutex.synchronize do
+        return if @schema_ready
+
+        connection.execute(create_table_sql)
+        @schema_ready = true
+      end
+    end
+
+    private
+
+    def connection
+      PartitionGardener.configuration.connection
+    end
+
+    def quoted_table(name)
+      connection.quote_table_name(name)
+    end
+
+    def load_sql(table_name)
+      <<~SQL
+        SELECT table_name, phase, plan_signature, staging_row_count
+        FROM #{quoted_table(TABLE_NAME)}
+        WHERE table_name = #{connection.quote(table_name)}
+      SQL
+    end
+
+    def save_sql(table_name, attributes)
+      <<~SQL
+        INSERT INTO #{quoted_table(TABLE_NAME)} (
+          table_name, phase, plan_signature, staging_row_count, updated_at
+        ) VALUES (
+          #{connection.quote(table_name)},
+          #{connection.quote(attributes.fetch(:phase))},
+          #{connection.quote(attributes.fetch(:plan_signature))},
+          #{attributes.fetch(:staging_row_count, 0).to_i},
+          NOW()
+        )
+        ON CONFLICT (table_name) DO UPDATE SET
+          phase = EXCLUDED.phase,
+          plan_signature = EXCLUDED.plan_signature,
+          staging_row_count = EXCLUDED.staging_row_count,
+          updated_at = NOW()
+      SQL
+    end
+
+    def clear_sql(table_name)
+      <<~SQL
+        DELETE FROM #{quoted_table(TABLE_NAME)}
+        WHERE table_name = #{connection.quote(table_name)}
+      SQL
+    end
+
+    def create_table_sql
+      <<~SQL
+        CREATE TABLE IF NOT EXISTS #{quoted_table(TABLE_NAME)} (
+          table_name text PRIMARY KEY,
+          phase text NOT NULL,
+          plan_signature text NOT NULL,
+          staging_row_count integer NOT NULL DEFAULT 0,
+          updated_at timestamptz NOT NULL DEFAULT NOW()
+        )
+      SQL
+    end
+  end
+
+  ActiveRecordRunRecordStore = SqlRunRecordStore
+end

data/lib/partition_gardener/stdlib_extensions.rb

@@ -0,0 +1,15 @@
+module PartitionGardener
+  module IntegerDurationMethods
+    def minutes
+      self * 60
+    end
+
+    def seconds
+      self
+    end
+  end
+end
+
+unless defined?(ActiveSupport::Duration)
+  Integer.include(PartitionGardener::IntegerDurationMethods)
+end

data/lib/partition_gardener/strategy/composite.rb

@@ -0,0 +1,27 @@
+module PartitionGardener
+  class CompositeMaintenance
+    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!
+      Templates.expand(@config).each do |branch_config|
+        ThreeAreaMaintenance.new(branch_config, job_class_name: @job_class_name, executor: @executor).run!
+      end
+    end
+
+    def split_future_month_from_current!(_identifier = nil)
+      run!
+    end
+
+    def split_pressured_future_month_partitions
+      run!
+    end
+
+    def collapse_low_volume_future_month_partitions
+      run!
+    end
+  end
+end

data/lib/partition_gardener/strategy/cursor_columns.rb

@@ -0,0 +1,18 @@
+module PartitionGardener
+  module Strategy
+    module CursorColumns
+      def cursor_columns
+        base_key = partition_key_base
+        conflict_columns = @config[:conflict_key] || [base_key, "id"]
+        ([base_key] + conflict_columns).uniq
+      end
+
+      def partition_key_base
+        column = @config[:partition_key_column]
+        return column.to_s if column.nil?
+
+        column.to_s.split("::").first.strip
+      end
+    end
+  end
+end

data/lib/partition_gardener/strategy/date_range.rb

@@ -0,0 +1,303 @@
+module PartitionGardener
+  module Strategy
+    class DateRange
+      include Naming
+      include RequiresDefaultPartition
+      include CursorColumns
+
+      DEFAULT_ACTIVE_MONTHS = 12
+      DEFAULT_ACTIVE_YEARS = 2
+
+      def initialize(config)
+        @config = config
+      end
+
+      def active_window
+        bucket = date_bucket
+        active_start = DateBucket.beginning_of_bucket(today, bucket)
+        active_span = active_span_for(bucket)
+        active_end = DateBucket.add_buckets(active_start, active_span, bucket)
+
+        {start: active_start, end: active_end}
+      end
+
+      def build_plan
+        window = active_window
+        heatmap = collect_heatmap(window)
+        hot_buckets = hot_buckets_in_window(heatmap, window)
+
+        segments = if year_bucket?
+          Layout::CalendarYear.build_segments(
+            config: @config,
+            active_start: window[:start],
+            active_end: window[:end],
+            hot_years: hot_buckets
+          )
+        else
+          Layout::SlidingWindow.build_segments(
+            config: @config,
+            active_start: window[:start],
+            active_end: window[:end],
+            hot_months: hot_buckets
+          )
+        end
+
+        Plan::Result.new(segments: segments, hot_buckets: hot_buckets)
+      end
+
+      def attached_tail_segments
+        window = active_window
+
+        Connection.attached_partitions(table_name).filter_map do |partition|
+          next if partition.default
+          next unless managed_tail_partition?(partition.name, window: window)
+
+          Plan::Segment.new(
+            name: partition.name,
+            range_start: partition.range_start,
+            range_end: partition.range_end,
+            kind: segment_kind_for(partition.name, window: window)
+          )
+        end
+      end
+
+      def collect_heatmap(window)
+        bucket_counts = Hash.new(0)
+        dedicated_partition_counts = {}
+        dedicated_partitions = hot_bucket_partitions_in_window(window).to_h
+
+        heatmap_source_partitions(window).each do |partition_name|
+          counts = counts_by_bucket_in_partition(partition_name)
+          counts.each { |bucket, count| bucket_counts[bucket] += count }
+
+          bucket = dedicated_partitions[partition_name]
+          dedicated_partition_counts[bucket] = counts.values.sum if bucket
+        end
+
+        {
+          bucket_counts: bucket_counts,
+          default_rows: default_row_count,
+          dedicated_partition_counts: dedicated_partition_counts
+        }
+      end
+
+      def hot_buckets_in_window(heatmap, window)
+        return [] if rolling_current_layout?
+
+        hot_buckets = Set.new
+
+        heatmap[:bucket_counts].each do |bucket, row_count|
+          next unless bucket_in_window?(bucket, window)
+          next unless row_count >= split_row_threshold
+
+          hot_buckets << bucket
+        end
+
+        dedicated_counts = heatmap.fetch(:dedicated_partition_counts, {})
+        hot_bucket_partitions_in_window(window).each do |_partition_name, bucket|
+          row_count = dedicated_counts[bucket]
+          next if row_count.nil?
+
+          if row_count >= split_row_threshold
+            hot_buckets << bucket
+          else
+            hot_buckets.delete(bucket)
+          end
+        end
+
+        hot_buckets.sort
+      end
+
+      def managed_tail_partition_names(window: active_window)
+        Connection.attached_partitions(table_name).filter_map do |partition|
+          next if partition.default
+          next unless managed_tail_partition?(partition.name, window: window)
+
+          partition.name
+        end
+      end
+
+      def tail_slot_name?(partition_name)
+        partition_name == current_partition_name(table_name) ||
+          partition_name == open_partition_name(table_name) ||
+          partition_name == future_partition_name(table_name) ||
+          partition_name.match?(/^#{Regexp.escape(table_name)}_open_\d+$/)
+      end
+
+      def current_and_future_where_condition(window: active_window)
+        partition_key_column = @config[:partition_key_column]
+        "#{partition_key_column} >= #{connection.quote(window[:start])}::date"
+      end
+
+      def bucket_where_condition(bucket)
+        partition_key_column = @config[:partition_key_column]
+        start_range = beginning_of_bucket(bucket)
+        end_range = end_of_bucket(bucket)
+        "#{partition_key_column} >= #{connection.quote(start_range)}::date AND #{partition_key_column} < #{connection.quote(end_range)}::date"
+      end
+
+      def archive_bucket?(bucket)
+        beginning_of_bucket(bucket) < active_window[:start]
+      end
+
+      def future_bucket?(bucket)
+        beginning_of_bucket(bucket) > beginning_of_bucket(today)
+      end
+
+      def archive_bucket_from_partition_name(partition_name)
+        DateBucket.archive_bucket_from_partition_name(table_name, partition_name, date_bucket)
+      end
+
+      def segment_for_values_clause(segment)
+        if segment.range_end == :max
+          "FROM ('#{segment.range_start}') TO (MAXVALUE)"
+        elsif segment.hash_partition?
+          "WITH (modulus #{segment.range_start[:modulus]}, remainder #{segment.range_start[:remainder]})"
+        else
+          "FROM ('#{segment.range_start}') TO ('#{segment.range_end}')"
+        end
+      end
+
+      def bucket_counts_in_partition(partition_name)
+        counts_by_bucket_in_partition(partition_name)
+      end
+
+      private
+
+      def table_name
+        @config[:table_name]
+      end
+
+      def connection
+        Connection.connection
+      end
+
+      def today
+        PartitionGardener.configuration.today
+      end
+
+      def date_bucket
+        DateBucket.normalize(@config.fetch(:bucket, :month))
+      end
+
+      def year_bucket?
+        date_bucket == :year
+      end
+
+      def rolling_current_layout?
+        @config.fetch(:layout, :sliding_window) == :rolling_current
+      end
+
+      def active_span_for(bucket)
+        active_key = DateBucket.active_key(bucket)
+        default = case bucket
+        when :year then DEFAULT_ACTIVE_YEARS
+        when :month then DEFAULT_ACTIVE_MONTHS
+        else DateBucket.default_active_span(bucket)
+        end
+        @config.fetch(active_key, default)
+      end
+
+      def split_row_threshold
+        @config.fetch(:split_row_threshold, FUTURE_MONTH_PARTITION_ROW_THRESHOLD)
+      end
+
+      def beginning_of_bucket(bucket)
+        DateBucket.beginning_of_bucket(bucket.to_date, date_bucket)
+      end
+
+      def end_of_bucket(bucket)
+        DateBucket.end_of_bucket(bucket.to_date, date_bucket)
+      end
+
+      def each_bucket_in_range(range_start, range_end_exclusive)
+        buckets = []
+        bucket = beginning_of_bucket(range_start)
+
+        while bucket < range_end_exclusive
+          buckets << bucket
+          bucket = DateBucket.next_bucket(bucket, date_bucket)
+        end
+
+        buckets
+      end
+
+      def bucket_in_window?(bucket, window)
+        return false if window[:start].nil? || window[:end].nil?
+
+        bucket >= window[:start] && bucket < window[:end]
+      end
+
+      def managed_tail_partition?(partition_name, window:)
+        return true if tail_slot_name?(partition_name)
+
+        bucket = archive_bucket_from_partition_name(partition_name)
+        bucket && bucket_in_window?(bucket, window)
+      end
+
+      def segment_kind_for(partition_name, window:)
+        return :future if partition_name == future_partition_name(table_name)
+        return :filler if tail_slot_name?(partition_name)
+
+        :hot_bucket
+      end
+
+      def heatmap_source_partitions(window)
+        names = tail_slot_names_for_heatmap.select { |name| Connection.partition_attached?(table_name, name) }
+        names.concat(hot_bucket_partitions_in_window(window).map(&:first))
+        names.uniq
+      end
+
+      def tail_slot_names_for_heatmap
+        [
+          current_partition_name(table_name),
+          open_partition_name(table_name),
+          future_partition_name(table_name)
+        ] + open_slot_partition_names
+      end
+
+      def open_slot_partition_names
+        Connection.attached_partitions(table_name).map(&:name).select do |partition_name|
+          partition_name.match?(/^#{Regexp.escape(table_name)}_open_\d+$/)
+        end
+      end
+
+      def hot_bucket_partitions_in_window(window)
+        Connection.attached_partitions(table_name).filter_map do |partition|
+          bucket = archive_bucket_from_partition_name(partition.name)
+          next unless bucket
+          next unless bucket_in_window?(bucket, window)
+
+          [partition.name, bucket]
+        end
+      end
+
+      def counts_by_bucket_in_partition(partition_name)
+        partition_key_column = @config[:partition_key_column]
+        trunc_unit = DateBucket.date_trunc_unit(date_bucket)
+        bucket_expression = if partition_key_column.include?("::")
+          "date_trunc('#{trunc_unit}', #{partition_key_column})::date"
+        else
+          "date_trunc('#{trunc_unit}', #{connection.quote_column_name(partition_key_column)})::date"
+        end
+
+        sql = <<~SQL
+          SELECT #{bucket_expression} AS bucket, COUNT(*)::int AS row_count
+          FROM #{Connection.quoted_table(partition_name)}
+          GROUP BY 1
+        SQL
+
+        connection.execute(sql).each_with_object({}) do |row, counts|
+          counts[Date.parse(row["bucket"].to_s)] = row["row_count"].to_i
+        end
+      end
+
+      def default_row_count
+        default_name = default_partition_name(table_name)
+        return 0 unless Connection.partition_exists?(default_name)
+
+        Connection.count_rows_in_partition_table(default_name)
+      end
+    end
+  end
+end