pgdice 0.1.0 → 0.2.0

data/lib/pgdice/configuration_file_loader.rb

@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+module PgDice
+  #  ConfigurationFileLoader is a class used to load the PgDice configuration file
+  class ConfigurationFileLoader
+    include PgDice::LogHelper
+    extend Forwardable
+
+    attr_reader :config
+
+    def_delegators :@config, :config_file, :logger
+
+    def initialize(config = PgDice::Configuration.new, opts = {})
+      @config = config
+      @file_validator = opts[:file_validator] ||= lambda do |config_file|
+        validate_file(config_file)
+      end
+      @config_loader = opts[:config_loader] ||= lambda do |file|
+        logger.debug { "Loading PgDice configuration file: '#{config_file}'" }
+        YAML.safe_load(ERB.new(IO.read(file)).result)
+      end
+      @file_loaded = opts[:file_loaded]
+    end
+
+    def load_file
+      return if @file_loaded
+
+      @file_loaded = true
+
+      @file_validator.call(config_file)
+
+      @config.approved_tables = @config_loader.call(config_file)
+                                              .fetch('approved_tables')
+                                              .reduce(tables(@config)) do |tables, hash|
+        tables << PgDice::Table.from_hash(hash)
+      end
+    end
+
+    def file_loaded?
+      @file_loaded
+    end
+
+    private
+
+    def validate_file(config_file)
+      if blank?(config_file)
+        raise PgDice::InvalidConfigurationError,
+              'Cannot read in nil configuration file! You must set config_file if you leave approved_tables nil!'
+      end
+
+      raise PgDice::MissingConfigurationFileError, config_file unless File.exist?(config_file)
+    end
+
+    def tables(config)
+      if config.approved_tables(eager_load: true).is_a?(PgDice::ApprovedTables)
+        return config.approved_tables(eager_load: true)
+      end
+
+      PgDice::ApprovedTables.new
+    end
+  end
+end

data/lib/pgdice/database_connection.rb

@@ -4,25 +4,33 @@
 module PgDice
   # Wrapper class around database connection handlers
   class DatabaseConnection
-    extend Forwardable
-    def_delegators :@configuration, :logger, :dry_run, :pg_connection
+    include PgDice::LogHelper
 
-    def initialize(configuration = PgDice::Configuration.new)
-      @configuration = configuration
+    attr_reader :logger, :query_executor, :dry_run
+
+    def initialize(logger:, query_executor:, dry_run: false)
+      @logger = logger
+      @dry_run = dry_run
+      @query_executor = query_executor
     end
 
     def execute(query)
+      query = squish(query)
+
+      if blank?(query) || dry_run
+        logger.debug { "DatabaseConnection skipping query. Query: '#{query}'. Dry run: #{dry_run}" }
+        return PgDice::PgResponse.new
+      end
+
       logger.debug { "DatabaseConnection to execute query: #{query}" }
-      if dry_run
-        PgDicePgResponse.new
-      else
-        pg_connection.exec(query)
+      PgDice::LogHelper.log_duration('Executing query', logger) do
+        query_executor.call(query)
       end
     end
   end
 
   # Null-object pattern for PG::Result since that object isn't straightforward to initialize
-  class PgDicePgResponse
+  class PgResponse
     def values
       []
     end

data/lib/pgdice/database_connection_factory.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+# Entry point for PartitionManager
+module PgDice
+  #  PartitionListerFactory is a class used to build PartitionListers
+  class DatabaseConnectionFactory
+    extend Forwardable
+
+    def_delegators :@configuration, :logger, :pg_connection, :dry_run
+
+    def initialize(configuration, opts = {})
+      @configuration = configuration
+      @query_executor = opts[:query_executor] ||= ->(query) { pg_connection.exec(query) }
+    end
+
+    def call
+      PgDice::DatabaseConnection.new(logger: logger, query_executor: @query_executor, dry_run: dry_run)
+    end
+  end
+end

data/lib/pgdice/date_helper.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module PgDice
+  # Helper used to manipulate date objects
+  module DateHelper
+    def pad_date(numbers)
+      return numbers if numbers.size == 8
+
+      case numbers.size
+      when 6
+        return numbers + '01'
+      when 4
+        return numbers + '0101'
+      else
+        raise ArgumentError, "Invalid date. Cannot parse date from #{numbers}"
+      end
+    end
+
+    def truncate_date(date, period)
+      case period
+      when 'year'
+        Date.parse("#{date.year}0101")
+      when 'month'
+        Date.parse("#{date.year}#{date.month}01")
+      when 'day'
+        date
+      else
+        raise ArgumentError, "Invalid date. Cannot parse date from #{date}"
+      end
+    end
+
+    def safe_date_builder(table_name)
+      matches = table_name.match(/\d+/)
+      raise ArgumentError, "Invalid date. Cannot parse date from #{table_name}" unless matches
+
+      Date.parse(pad_date(matches[0]))
+    end
+  end
+end

data/lib/pgdice/error.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+module PgDice
+  # PgDice parent error class
+  class Error < StandardError
+  end
+
+  # Error thrown when PgSlice returns an error code
+  class PgSliceError < Error
+  end
+
+  # Generic validation error
+  class ValidationError < Error
+  end
+
+  # Error thrown if a user tries to operate on a table that is not in the ApprovedTables object.
+  class IllegalTableError < ValidationError
+  end
+
+  # Error thrown when a user attempts to manipulate partitions on a table that is not partitioned
+  class TableNotPartitionedError < Error
+  end
+
+  # Generic error for table counts
+  class InsufficientTablesError < Error
+    def initialize(direction, table_name, expected, period, found_count)
+      super("Insufficient #{direction} tables exist for table: #{table_name}. "\
+      "Expected: #{expected} having period of: #{period} but found: #{found_count}")
+    end
+  end
+
+  # Error thrown when the count of future tables is less than the expected amount
+  class InsufficientFutureTablesError < InsufficientTablesError
+    def initialize(table_name, expected, period, found_count)
+      super('future', table_name, expected, period, found_count)
+    end
+  end
+
+  # Error thrown when the count of past tables is less than the expected amount
+  class InsufficientPastTablesError < InsufficientTablesError
+    def initialize(table_name, expected, period, found_count)
+      super('past', table_name, expected, period, found_count)
+    end
+  end
+
+  # Generic configuration error
+  class ConfigurationError < Error
+  end
+
+  # Error thrown if you call a method that requires configuration first
+  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
+
+  # Error thrown if you provide bad data in a configuration
+  class InvalidConfigurationError < ConfigurationError
+    def initialize(message)
+      super("PgDice is not configured properly. #{message}")
+    end
+  end
+
+  # Error thrown if the config file specified does not exist.
+  class MissingConfigurationFileError < ConfigurationError
+    def initialize(config_file)
+      super("File: '#{config_file}' could not be found or does not exist. Is this the correct file path?")
+    end
+  end
+end

data/lib/pgdice/log_helper.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module PgDice
+  # LogHelper provides a convenient wrapper block to log out the duration of an operation
+  module LogHelper
+    def blank?(string)
+      string.nil? || string.empty?
+    end
+
+    def squish(string)
+      string ||= ''
+      string.gsub(/\s+/, ' ')
+    end
+
+    class << self
+      # If you want to pass the the result of your block into the message you can use '{}' and it will be
+      # substituted with the result of your block.
+      def log_duration(message, logger, options = {})
+        logger.error { 'log_duration called without a block. Cannot time the duration of nothing.' } unless block_given?
+        time_start = Time.now.utc
+        result = yield
+        time_end = Time.now.utc
+
+        formatted_message = format_message(time_end, time_start, message, result)
+        logger.public_send(options[:log_level] || :debug) { formatted_message }
+        result
+      end
+
+      private
+
+      def format_message(time_end, time_start, message, result)
+        message = message.sub(/{}/, result.to_s)
+        "#{message} took: #{format('%.02f', (time_end - time_start))} seconds."
+      end
+    end
+  end
+end

data/lib/pgdice/partition_dropper.rb

@@ -0,0 +1,34 @@
+# 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 PartitionDropper
+    attr_reader :logger, :query_executor
+
+    def initialize(logger:, query_executor:)
+      @logger = logger
+      @query_executor = query_executor
+    end
+
+    def call(old_partitions)
+      logger.info { "Partitions to be deleted are: #{old_partitions}" }
+
+      query_executor.call(generate_drop_sql(old_partitions))
+
+      old_partitions
+    end
+
+    private
+
+    def generate_drop_sql(old_partitions)
+      return if old_partitions.size.zero?
+
+      sql_query = old_partitions.reduce("BEGIN;\n") do |sql, table_name|
+        sql + "DROP TABLE IF EXISTS #{table_name} CASCADE;\n"
+      end
+      sql_query + 'COMMIT;'
+    end
+  end
+end

data/lib/pgdice/partition_dropper_factory.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+# Entry point for PartitionManager
+module PgDice
+  #  PartitionListerFactory is a class used to build PartitionListers
+  class PartitionDropperFactory
+    extend Forwardable
+
+    def_delegators :@configuration, :logger, :database_connection
+
+    def initialize(configuration, opts = {})
+      @configuration = configuration
+      @query_executor = opts[:query_executor] ||= ->(sql) { database_connection.execute(sql) }
+    end
+
+    def call
+      PgDice::PartitionDropper.new(logger: logger, query_executor: @query_executor)
+    end
+  end
+end

data/lib/pgdice/partition_helper.rb

@@ -4,49 +4,37 @@
 module PgDice
   # Helps do high-level tasks like getting tables partitioned
   class PartitionHelper
-    extend Forwardable
-    def_delegators :@configuration, :logger
+    attr_reader :logger, :approved_tables, :validation, :pg_slice_manager
 
-    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)
+    def initialize(logger:, approved_tables:, validation:, pg_slice_manager:)
+      @logger = logger
+      @validation = validation
+      @approved_tables = approved_tables
+      @pg_slice_manager = pg_slice_manager
     end
 
-    def partition_table!(params = {})
-      params[:column_name] ||= 'created_at'
-      params[:period] ||= :day
-
-      validation_helper.validate_parameters(params)
+    def partition_table(table_name, params = {})
+      table = approved_tables.fetch(table_name)
+      all_params = table.smash(params)
+      validation.validate_parameters(all_params)
 
-      logger.info { "Preparing database with params: #{params}" }
+      logger.info { "Preparing database for table: #{table}. Using parameters: #{all_params}" }
 
-      prep_and_fill(params)
-      swap_and_fill(params)
+      prep_and_fill(all_params)
+      swap_and_fill(all_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}" }
+    def undo_partitioning!(table_name)
+      approved_tables.fetch(table_name)
+      logger.info { "Undoing partitioning for table: #{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)
+    def undo_partitioning(table_name)
+      undo_partitioning!(table_name)
     rescue PgDice::PgSliceError => error
       logger.error { "Rescued PgSliceError: #{error}" }
       false

data/lib/pgdice/partition_helper_factory.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+# Entry point for PartitionManager
+module PgDice
+  #  PartitionManagerFactory is a class used to build PartitionManagers
+  class PartitionHelperFactory
+    extend Forwardable
+
+    def_delegators :@configuration, :logger, :approved_tables
+
+    def initialize(configuration, opts = {})
+      @configuration = configuration
+      @validation_factory = opts[:validation_factory] ||= PgDice::ValidationFactory.new(configuration)
+      @pg_slice_manager_factory = opts[:pg_slice_manager_factory] ||= PgDice::PgSliceManagerFactory.new(configuration)
+    end
+
+    def call
+      PgDice::PartitionHelper.new(logger: logger,
+                                  approved_tables: approved_tables,
+                                  validation: @validation_factory.call,
+                                  pg_slice_manager: @pg_slice_manager_factory.call)
+    end
+  end
+end

data/lib/pgdice/partition_lister.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+# Entry point for PartitionManager
+module PgDice
+  #  PartitionLister is used to list partitions
+  class PartitionLister
+    attr_reader :query_executor
+
+    def initialize(query_executor:)
+      @query_executor = query_executor
+    end
+
+    def call(all_params)
+      sql = build_partition_table_fetch_sql(all_params)
+      query_executor.call(sql)
+    end
+
+    private
+
+    def build_partition_table_fetch_sql(params = {})
+      schema = params.fetch(:schema, 'public')
+      base_table_name = params.fetch(:table_name)
+
+      <<~SQL
+        SELECT tablename
+        FROM pg_tables
+        WHERE schemaname = '#{schema}'
+          AND tablename ~ '^#{base_table_name}_\\d+$'
+        ORDER BY tablename;
+      SQL
+    end
+  end
+end

data/lib/pgdice/partition_lister_factory.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+# Entry point for PartitionManager
+module PgDice
+  #  PartitionListerFactory is a class used to build PartitionListers
+  class PartitionListerFactory
+    extend Forwardable
+
+    def_delegators :@configuration, :database_connection
+
+    def initialize(configuration, opts = {})
+      @configuration = configuration
+      @query_executor = opts[:query_executor] ||= lambda do |sql|
+        database_connection.execute(sql).values.flatten
+      end
+    end
+
+    def call
+      PgDice::PartitionLister.new(query_executor: @query_executor)
+    end
+  end
+end

data/lib/pgdice/partition_manager.rb

@@ -4,83 +4,87 @@
 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)
+    include PgDice::TableFinder
+
+    attr_reader :logger, :batch_size, :validation, :approved_tables, :partition_adder,
+                :partition_lister, :partition_dropper, :current_date_provider
+
+    def initialize(opts = {})
+      @logger = opts.fetch(:logger)
+      @batch_size = opts.fetch(:batch_size)
+      @validation = opts.fetch(:validation)
+      @approved_tables = opts.fetch(:approved_tables)
+      @partition_adder = opts.fetch(:partition_adder)
+      @partition_lister = opts.fetch(:partition_lister)
+      @partition_dropper = opts.fetch(:partition_dropper)
+      @current_date_provider = opts.fetch(:current_date_provider, proc { Time.now.utc.to_date })
     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)
+    def add_new_partitions(table_name, params = {})
+      all_params = approved_tables.smash(table_name, params)
+      logger.debug { "add_new_partitions has been called with params: #{all_params}" }
+      validation.validate_parameters(all_params)
+      partition_adder.call(all_params)
     end
 
-    def drop_old_partitions(params = {})
-      logger.info { "drop_old_partitions has been called with params: #{params}" }
+    def drop_old_partitions(table_name, params = {})
+      all_params = approved_tables.smash(table_name, params)
+      all_params[:older_than] = current_date_provider.call
+      logger.debug { "drop_old_partitions has been called with params: #{all_params}" }
 
-      validation.validate_parameters(params)
-      old_partitions = list_old_partitions(params)
-      logger.warn { "Partitions to be deleted are: #{old_partitions}" }
+      validation.validate_parameters(all_params)
+      drop_partitions(all_params)
+    end
 
-      old_partitions.each do |old_partition|
-        @configuration.table_dropper.call(old_partition, logger)
-      end
-      old_partitions
+    # Grabs only tables that start with the base_table_name and end in numbers
+    def list_partitions(table_name, params = {})
+      all_params = approved_tables.smash(table_name, params)
+      validation.validate_parameters(all_params)
+      partitions(all_params)
     end
 
-    def list_old_partitions(params = {})
-      params[:older_than] ||= older_than
-      logger.info { "Listing old partitions with params: #{params}" }
+    def list_droppable_partitions(table_name, params = {})
+      all_params = approved_tables.smash(table_name, params)
+      validation.validate_parameters(all_params)
+      droppable_partitions(all_params)
+    end
 
-      validation.validate_parameters(params)
+    def list_batched_droppable_partitions(table_name, params = {})
+      all_params = approved_tables.smash(table_name, params)
+      validation.validate_parameters(all_params)
+      droppable_tables = batched_droppable_partitions(all_params)
+      logger.debug { "Batched partitions eligible for dropping are: #{droppable_tables}" }
+      droppable_tables
+    end
 
-      partition_tables = fetch_partition_tables(params)
+    private
 
-      filter_partitions(partition_tables, params[:table_name], params[:older_than])
+    def partitions(all_params)
+      logger.info { "Fetching partition tables with params: #{all_params}" }
+      partition_lister.call(all_params)
     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}" }
+    def droppable_partitions(all_params)
+      older_than = current_date_provider.call
+      minimum_tables = all_params.fetch(:past)
+      period = all_params.fetch(:period)
 
-      sql = build_partition_table_fetch_sql(params)
+      eligible_partitions = partitions(all_params)
 
-      partition_tables = database_connection.execute(sql).values.flatten
-      logger.debug { "Table: #{schema}.#{params[:table_name]} has partition_tables: #{partition_tables}" }
-      partition_tables
+      droppable_tables = find_droppable_partitions(eligible_partitions, older_than, minimum_tables, period)
+      logger.debug { "Partitions eligible for dropping older than: #{older_than} are: #{droppable_tables}" }
+      droppable_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
+    def batched_droppable_partitions(all_params)
+      max_tables_to_drop_at_once = all_params.fetch(:batch_size, batch_size)
+      selected_partitions = droppable_partitions(all_params)
+      batched_tables(selected_partitions, max_tables_to_drop_at_once)
     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
+    def drop_partitions(all_params)
+      old_partitions = batched_droppable_partitions(all_params)
+      partition_dropper.call(old_partitions)
     end
   end
 end