pgdice 0.1.0

data/lib/pgdice/configuration.rb

@@ -0,0 +1,122 @@
+# frozen_string_literal: true
+
+# Entry point for configuration
+module PgDice
+  class << self
+    attr_accessor :configuration
+
+    def configure
+      self.configuration ||= PgDice::Configuration.new
+      yield(configuration)
+    end
+  end
+
+  # Configuration class which holds all configurable values
+  class Configuration
+    def self.days_ago(days)
+      Time.now.utc - days * 24 * 60 * 60
+    end
+
+    VALUES = { logger: Logger.new(STDOUT),
+               database_url: nil,
+               additional_validators: [],
+               approved_tables: [],
+               older_than: PgDice::Configuration.days_ago(90),
+               dry_run: false,
+               table_drop_batch_size: 7 }.freeze
+
+    attr_writer :logger,
+                :database_url,
+                :additional_validators,
+                :approved_tables,
+                :older_than,
+                :dry_run,
+                :table_drop_batch_size,
+                :database_connection,
+                :pg_connection
+
+    attr_accessor :table_dropper,
+                  :pg_slice_manager,
+                  :partition_manager,
+                  :partition_helper
+
+    def initialize(existing_configuration = nil)
+      VALUES.each do |key, value|
+        initialize_value(key, value, existing_configuration)
+      end
+      initialize_objects
+    end
+
+    def logger
+      return @logger unless @logger.nil?
+
+      raise PgDice::InvalidConfigurationError, 'logger must be present!'
+    end
+
+    def database_url
+      return @database_url unless @database_url.nil?
+
+      raise PgDice::InvalidConfigurationError, 'database_url must be present!'
+    end
+
+    def database_connection
+      return @database_connection unless @database_connection.nil?
+
+      raise PgDice::InvalidConfigurationError, 'database_connection must be present!'
+    end
+
+    def additional_validators
+      return @additional_validators if @additional_validators.is_a?(Array)
+
+      raise PgDice::InvalidConfigurationError, 'additional_validators must be an Array!'
+    end
+
+    def approved_tables
+      return @approved_tables if @approved_tables.is_a?(Array)
+
+      raise PgDice::InvalidConfigurationError, 'approved_tables must be an Array of strings!'
+    end
+
+    def older_than
+      return @older_than if @older_than.is_a?(Time)
+
+      raise PgDice::InvalidConfigurationError, 'older_than must be a Time!'
+    end
+
+    def dry_run
+      return @dry_run if [true, false].include?(@dry_run)
+
+      raise PgDice::InvalidConfigurationError, 'dry_run must be either true or false!'
+    end
+
+    def table_drop_batch_size
+      return @table_drop_batch_size.to_i if @table_drop_batch_size.to_i >= 0
+
+      raise PgDice::InvalidConfigurationError, 'table_drop_batch_size must be a non-negative Integer!'
+    end
+
+    # Lazily initialized
+    def pg_connection
+      @pg_connection ||= PG::Connection.new(database_url)
+      return @pg_connection if @pg_connection.respond_to?(:exec)
+
+      raise PgDice::InvalidConfigurationError, 'pg_connection must be present!'
+    end
+
+    def deep_clone
+      PgDice::Configuration.new(self)
+    end
+
+    private
+
+    def initialize_value(variable_name, default_value, existing_configuration)
+      instance_variable_set("@#{variable_name}", existing_configuration&.send(variable_name)&.clone || default_value)
+    end
+
+    def initialize_objects
+      @database_connection = PgDice::DatabaseConnection.new(self)
+      @partition_manager = PgDice::PartitionManager.new(self)
+      @table_dropper = PgDice::TableDropper.new(self)
+    end
+  end
+end

data/lib/pgdice/database_connection.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+# Entry point for DatabaseConnection
+module PgDice
+  # Wrapper class around database connection handlers
+  class DatabaseConnection
+    extend Forwardable
+    def_delegators :@configuration, :logger, :dry_run, :pg_connection
+
+    def initialize(configuration = PgDice::Configuration.new)
+      @configuration = configuration
+    end
+
+    def execute(query)
+      logger.debug { "DatabaseConnection to execute query: #{query}" }
+      if dry_run
+        PgDicePgResponse.new
+      else
+        pg_connection.exec(query)
+      end
+    end
+  end
+
+  # Null-object pattern for PG::Result since that object isn't straightforward to initialize
+  class PgDicePgResponse
+    def values
+      []
+    end
+  end
+end

data/lib/pgdice/partition_helper.rb

@@ -0,0 +1,69 @@
+# frozen_string_literal: true
+
+# Entry point for PartitionHelper
+module PgDice
+  # Helps do high-level tasks like getting tables partitioned
+  class PartitionHelper
+    extend Forwardable
+    def_delegators :@configuration, :logger
+
+    attr_reader :pg_slice_manager, :validation_helper
+
+    def initialize(configuration = PgDice::Configuration.new)
+      @configuration = configuration
+      @pg_slice_manager = PgDice::PgSliceManager.new(configuration)
+      @validation_helper = PgDice::Validation.new(configuration)
+    end
+
+    def partition_table!(params = {})
+      params[:column_name] ||= 'created_at'
+      params[:period] ||= :day
+
+      validation_helper.validate_parameters(params)
+
+      logger.info { "Preparing database with params: #{params}" }
+
+      prep_and_fill(params)
+      swap_and_fill(params)
+    end
+
+    def undo_partitioning!(params = {})
+      table_name = params.fetch(:table_name)
+
+      validation_helper.validate_parameters(params)
+      logger.info { "Cleaning up database with params: #{table_name}" }
+
+      pg_slice_manager.analyze(table_name: table_name, swapped: true)
+      pg_slice_manager.unswap!(table_name: table_name)
+      pg_slice_manager.unprep!(table_name: table_name)
+    end
+
+    def partition_table(params = {})
+      partition_table!(params)
+    rescue PgDice::PgSliceError => error
+      logger.error { "Rescued PgSliceError: #{error}" }
+      false
+    end
+
+    def undo_partitioning(params = {})
+      undo_partitioning!(params)
+    rescue PgDice::PgSliceError => error
+      logger.error { "Rescued PgSliceError: #{error}" }
+      false
+    end
+
+    private
+
+    def prep_and_fill(params)
+      pg_slice_manager.prep(params)
+      pg_slice_manager.add_partitions(params.merge!(intermediate: true))
+      pg_slice_manager.fill(params) if params[:fill]
+    end
+
+    def swap_and_fill(params)
+      pg_slice_manager.analyze(params)
+      pg_slice_manager.swap(params)
+      pg_slice_manager.fill(params.merge!(swapped: true)) if params[:fill]
+    end
+  end
+end

data/lib/pgdice/partition_manager.rb

@@ -0,0 +1,86 @@
+# frozen_string_literal: true
+
+# Entry point for PartitionManager
+module PgDice
+  #  PartitionManager is a class used to fulfill high-level tasks for partitioning
+  class PartitionManager
+    extend Forwardable
+    def_delegators :@configuration, :logger, :older_than, :table_drop_batch_size
+
+    attr_reader :validation, :pg_slice_manager, :database_connection
+
+    def initialize(configuration = PgDice::Configuration.new)
+      @configuration = configuration
+      @validation = PgDice::Validation.new(configuration)
+      @pg_slice_manager = PgDice::PgSliceManager.new(configuration)
+      @database_connection = PgDice::DatabaseConnection.new(configuration)
+    end
+
+    def add_new_partitions(params = {})
+      logger.info { "add_new_partitions has been called with params: #{params}" }
+
+      validation.validate_parameters(params)
+      pg_slice_manager.add_partitions(params)
+    end
+
+    def drop_old_partitions(params = {})
+      logger.info { "drop_old_partitions has been called with params: #{params}" }
+
+      validation.validate_parameters(params)
+      old_partitions = list_old_partitions(params)
+      logger.warn { "Partitions to be deleted are: #{old_partitions}" }
+
+      old_partitions.each do |old_partition|
+        @configuration.table_dropper.call(old_partition, logger)
+      end
+      old_partitions
+    end
+
+    def list_old_partitions(params = {})
+      params[:older_than] ||= older_than
+      logger.info { "Listing old partitions with params: #{params}" }
+
+      validation.validate_parameters(params)
+
+      partition_tables = fetch_partition_tables(params)
+
+      filter_partitions(partition_tables, params[:table_name], params[:older_than])
+    end
+
+    # Grabs only tables that start with the base_table_name and end in numbers
+    def fetch_partition_tables(params = {})
+      schema = params[:schema] ||= 'public'
+      logger.info { "Fetching partition tables with params: #{params}" }
+
+      sql = build_partition_table_fetch_sql(params)
+
+      partition_tables = database_connection.execute(sql).values.flatten
+      logger.debug { "Table: #{schema}.#{params[:table_name]} has partition_tables: #{partition_tables}" }
+      partition_tables
+    end
+
+    private
+
+    def filter_partitions(partition_tables, base_table_name, partitions_older_than_time)
+      partition_tables.select do |partition_name|
+        partition_created_at_date = Date.parse(partition_name.gsub(/#{base_table_name}_/, '')).to_time
+        partition_created_at_date < partitions_older_than_time
+      end
+    end
+
+    def build_partition_table_fetch_sql(params = {})
+      schema = params.fetch(:schema)
+      base_table_name = params.fetch(:table_name)
+      limit = params.fetch(:limit, table_drop_batch_size)
+
+      <<~SQL
+        SELECT tablename
+        FROM pg_tables
+        WHERE schemaname = '#{schema}'
+          AND tablename ~ '^#{base_table_name}_\\d+$'
+        ORDER BY tablename
+        LIMIT #{limit}
+      SQL
+    end
+  end
+end

data/lib/pgdice/pg_slice_manager.rb

@@ -0,0 +1,119 @@
+# frozen_string_literal: true
+
+# Entry point for PgSliceManager
+module PgDice
+  # PgSliceManager is a wrapper around PgSlice
+  class PgSliceManager
+    extend Forwardable
+    def_delegators :@configuration, :logger, :database_url, :dry_run
+
+    def initialize(configuration = PgDice::Configuration.new)
+      @configuration = configuration
+    end
+
+    def prep(params = {})
+      table_name = params.fetch(:table_name)
+      column_name = params.fetch(:column_name)
+      period = params.fetch(:period)
+      run_pgslice("prep #{table_name} #{column_name} #{period}")
+    end
+
+    def fill(params = {})
+      table_name = params.fetch(:table_name)
+      swapped = params.fetch(:swapped, '')
+      swapped = '--swapped' if swapped.to_s.casecmp('true').zero?
+
+      run_pgslice("fill #{table_name} #{swapped}")
+    end
+
+    def analyze(params = {})
+      table_name = params.fetch(:table_name)
+      swapped = params.fetch(:swapped, '')
+      swapped = '--swapped' if swapped.to_s.casecmp('true').zero?
+
+      run_pgslice("analyze #{table_name} #{swapped}")
+    end
+
+    def swap(params = {})
+      table_name = params.fetch(:table_name)
+      run_pgslice("swap #{table_name}")
+    end
+
+    def add_partitions(params = {})
+      table_name = params.fetch(:table_name)
+      future_tables = params.fetch(:future, nil)
+      future_tables = "--future #{Integer(future_tables)}" if future_tables
+
+      past_tables = params.fetch(:past, nil)
+      past_tables = "--past #{Integer(past_tables)}" if past_tables
+
+      intermediate = params.fetch(:intermediate, nil)
+      intermediate = '--intermediate' if intermediate.to_s.casecmp('true').zero?
+
+      run_pgslice("add_partitions #{table_name} #{intermediate} #{future_tables} #{past_tables}")
+    end
+
+    def unprep!(params = {})
+      table_name = params.fetch(:table_name)
+
+      run_pgslice("unprep #{table_name}")
+    end
+
+    def unswap!(params = {})
+      table_name = params.fetch(:table_name)
+
+      run_pgslice("unswap #{table_name}")
+    end
+
+    def unprep(params = {})
+      table_name = params.fetch(:table_name)
+
+      run_pgslice("unprep #{table_name}")
+    rescue PgSliceError => error
+      logger.error { "Rescued PgSliceError: #{error}" }
+      false
+    end
+
+    def unswap(params = {})
+      table_name = params.fetch(:table_name)
+
+      run_pgslice("unswap #{table_name}")
+    rescue PgSliceError => error
+      logger.error { "Rescued PgSliceError: #{error}" }
+      false
+    end
+
+    private
+
+    def run_pgslice(argument_string)
+      parameters = build_pg_slice_command(argument_string)
+
+      stdout, stderr, status = Open3.capture3(parameters)
+      log_result(stdout, stderr, status)
+
+      if status.exitstatus.to_i.positive?
+        raise PgDice::PgSliceError,
+              "pgslice with arguments: '#{argument_string}' failed with status: '#{status.exitstatus}' "\
+                 "STDOUT: '#{stdout}' STDERR: '#{stderr}'"
+      end
+      true
+    end
+
+    def build_pg_slice_command(argument_string)
+      argument_string = argument_string.strip
+      logger.info { "Running pgslice command: '#{argument_string}'" }
+      $stdout.flush
+      $stderr.flush
+      command = "pgslice #{argument_string} "
+      command += '--dry-run true ' if dry_run
+      command += "--url #{database_url}"
+      command
+    end
+
+    def log_result(stdout, stderr, status)
+      logger.debug { "pgslice STDERR: #{stderr}" } if stderr
+      logger.debug { "pgslice STDOUT: #{stdout}" } if stdout
+      logger.debug { "pgslice exit status: #{status.exitstatus}" } if status
+    end
+  end
+end

data/lib/pgdice/table_dropper.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+# Entry point for TableDropperHelper
+module PgDice
+  # Simple class used to provide a mechanism that users can hook into if they want to override this
+  # default behavior for dropping a table.
+  class TableDropper
+    def initialize(configuration = PgDice::Configuration.new)
+      @configuration = configuration
+    end
+
+    def call(table_to_drop, _logger)
+      @configuration.database_connection.execute(drop_partition(table_to_drop))
+    end
+
+    private
+
+    def drop_partition(table_name)
+      <<~SQL
+        BEGIN;
+          DROP TABLE IF EXISTS #{table_name} CASCADE;
+        COMMIT;
+      SQL
+    end
+  end
+end

data/lib/pgdice/validation.rb

@@ -0,0 +1,132 @@
+# frozen_string_literal: true
+
+module PgDice
+  # Collection of utilities that provide ways for users to ensure things are working properly
+  class Validation
+    extend Forwardable
+    def_delegators :@configuration, :logger, :additional_validators, :database_connection, :approved_tables
+
+    def initialize(configuration = PgDice::Configuration.new)
+      @configuration = configuration
+    end
+
+    def assert_tables(params = {})
+      unless params[:future] || params[:past]
+        raise ArgumentError, 'You must provide either a future or past number of tables to assert on.'
+      end
+
+      table_name = validate_table_name(params)
+      period = resolve_period(params)
+
+      assert_future_tables(table_name, params[:future], period) if params[:future]
+      assert_past_tables(table_name, params[:past], period) if params[:past]
+      true
+    end
+
+    def validate_parameters(params)
+      validate_table_name(params)
+      validate_period(params)
+
+      run_additional_validators(params)
+      true
+    end
+
+    private
+
+    def resolve_period(params)
+      period = validate_period(params) || fetch_period_from_table_comment(params.fetch(:table_name))
+
+      # If the user doesn't supply a period and we fail to find one on the table then it's a pretty good bet
+      # this table is not partitioned at all.
+      unless period
+        raise TableNotPartitionedError,
+              "Table: #{params.fetch(:table_name)} is not partitioned! Cannot validate partitions that don't exist!"
+      end
+      period
+    end
+
+    def run_additional_validators(params)
+      return true if additional_validators.all? { |validator| validator.call(params, logger) }
+
+      raise PgDice::CustomValidationError.new(params, additional_validators)
+    rescue StandardError => error
+      raise PgDice::CustomValidationError.new(params, additional_validators, error)
+    end
+
+    def validate_table_name(params)
+      table_name = params.fetch(:table_name)
+      unless approved_tables.include?(table_name)
+        raise PgDice::IllegalTableError, "Table: #{table_name} is not in the list of approved tables!"
+      end
+
+      table_name
+    end
+
+    def validate_period(params)
+      return unless params[:period]
+
+      unless PgDice::SUPPORTED_PERIODS.include?(params[:period])
+        raise ArgumentError,
+              "Period must be one of: #{PgDice::SUPPORTED_PERIODS.keys}. Value: #{params[:period]} is not valid."
+      end
+
+      params[:period].to_sym
+    end
+
+    def assert_future_tables(table_name, future, period)
+      sql = build_assert_sql(table_name, future, period, :future)
+
+      response = database_connection.execute(sql)
+
+      return true if response.values.size == 1
+
+      raise PgDice::InsufficientFutureTablesError.new(table_name, future, period)
+    end
+
+    def assert_past_tables(table_name, past, period)
+      sql = build_assert_sql(table_name, past, period, :past)
+
+      response = database_connection.execute(sql)
+
+      return true if response.values.size == 1
+
+      raise PgDice::InsufficientPastTablesError.new(table_name, past, period)
+    end
+
+    def build_assert_sql(table_name, table_count, period, direction)
+      add_or_subtract = { future: '+', past: '-' }.fetch(direction, '+')
+      <<~SQL
+        SELECT 1
+        FROM pg_catalog.pg_class pg_class
+        INNER JOIN pg_catalog.pg_namespace pg_namespace ON pg_namespace.oid = pg_class.relnamespace
+        WHERE pg_class.relkind = 'r'
+          AND pg_namespace.nspname = 'public'
+          AND pg_class.relname = '#{table_name}_' || to_char(NOW()
+            #{add_or_subtract} INTERVAL '#{table_count} #{period}', '#{SUPPORTED_PERIODS[period.to_sym]}')
+      SQL
+    end
+
+    def fetch_period_from_table_comment(table_name)
+      sql = build_table_comment_sql(table_name, 'public')
+      values = database_connection.execute(sql).values.flatten.compact
+      convert_comment_to_hash(values.first)[:period]
+    end
+
+    def convert_comment_to_hash(comment)
+      return {} unless comment
+
+      partition_template = {}
+      comment.split(',').each do |key_value_pair|
+        key, value = key_value_pair.split(':')
+        partition_template[key.to_sym] = value
+      end
+      partition_template
+    end
+
+    def build_table_comment_sql(table_name, schema)
+      <<~SQL
+        SELECT obj_description('#{schema}.#{table_name}'::REGCLASS) AS comment
+      SQL
+    end
+  end
+end

data/lib/pgdice/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module PgDice
+  VERSION = '0.1.0'
+end

data/lib/pgdice.rb

@@ -0,0 +1,109 @@
+# frozen_string_literal: true
+
+require 'pg'
+require 'open3'
+require 'logger'
+require 'pgslice'
+require 'pgdice/version'
+require 'pgdice/validation'
+require 'pgdice/table_dropper'
+require 'pgdice/configuration'
+require 'pgdice/pg_slice_manager'
+require 'pgdice/partition_manager'
+require 'pgdice/partition_helper'
+require 'pgdice/database_connection'
+
+# This is a stupid comment
+module PgDice
+  class Error < StandardError
+  end
+  class PgSliceError < Error
+  end
+  class ValidationError < Error
+  end
+  class IllegalTableError < ValidationError
+  end
+  class TableNotPartitionedError < Error
+  end
+
+  # Rubocop is stupid
+  class InsufficientTablesError < Error
+    def initialize(direction, table_name, table_count, period)
+      super("Insufficient #{direction} tables exist for table: #{table_name}. "\
+        "Expected: #{table_count} having period of: #{period}")
+    end
+  end
+
+  # Rubocop is stupid
+  class InsufficientFutureTablesError < InsufficientTablesError
+    def initialize(table_name, table_count, period)
+      super('future', table_name, table_count, period)
+    end
+  end
+
+  # Rubocop is stupid
+  class InsufficientPastTablesError < InsufficientTablesError
+    def initialize(table_name, table_count, period)
+      super('past', table_name, table_count, period)
+    end
+  end
+
+  # Rubocop is stupid
+  class CustomValidationError < ValidationError
+    def initialize(params, validators, error = nil)
+      error_message = "Custom validation failed with params: #{params}. "
+      error_message += "Caused by error: #{error} " if error
+      error_message += "Validators: #{validators.map { |validator| source_location(validator) }.flatten}"
+      super(error_message)
+    end
+
+    private
+
+    # Helps users know what went wrong in their custom validators
+    def source_location(proc)
+      return proc.source_location if proc.respond_to?(:source_location)
+
+      proc.to_s
+    end
+  end
+
+  class ConfigurationError < Error
+  end
+
+  # Rubocop is stupid
+  class NotConfiguredError < ConfigurationError
+    def initialize(method_name)
+      super("Cannot use #{method_name} before PgDice has been configured! "\
+          'See README.md for configuration help.')
+    end
+  end
+
+  # Rubocop is stupid
+  class InvalidConfigurationError < ConfigurationError
+    def initialize(message)
+      super("PgDice is not configured properly. #{message}")
+    end
+  end
+
+  SUPPORTED_PERIODS = { day: 'YYYYMMDD', month: 'YYYYMM', year: 'YYYY' }.freeze
+
+  class << self
+    def partition_manager
+      raise PgDice::NotConfiguredError, 'partition_manager' unless configuration
+
+      @partition_manager ||= PgDice::PartitionManager.new(configuration)
+    end
+
+    def partition_helper
+      raise PgDice::NotConfiguredError, 'partition_helper' unless configuration
+
+      @partition_helper ||= PgDice::PartitionHelper.new(configuration)
+    end
+
+    def validation
+      raise PgDice::NotConfiguredError, 'validation' unless configuration
+
+      @validation ||= PgDice::Validation.new(configuration)
+    end
+  end
+end