pg_ha_migrations 1.8.0 → 2.1.0

data/lib/pg_ha_migrations/partman_config.rb

@@ -1,11 +1,73 @@
 # This is an internal class that is not meant to be used directly
 class PgHaMigrations::PartmanConfig < ActiveRecord::Base
+  SUPPORTED_PARTITION_TYPES = %w[native range]
+
+  delegate :connection, to: :class
+
   self.primary_key = :parent_table
 
-  # This method is called by unsafe_partman_update_config to set the fully
-  # qualified table name, as partman is often installed in a schema that
-  # is not included the application's search path
-  def self.schema=(schema)
-    self.table_name = "#{schema}.part_config"
+  def self.find(parent_table, partman_extension:)
+    unless partman_extension.installed?
+      raise PgHaMigrations::MissingExtensionError, "The pg_partman extension is not installed"
+    end
+
+    self.table_name = "#{partman_extension.quoted_schema}.part_config"
+
+    super(parent_table)
+  end
+
+  # The actual column type is TEXT and the value is determined by the
+  # intervalstyle in Postgres at the time create_parent is called.
+  # Rails hard codes this config when it builds connections for ease
+  # of parsing by ActiveSupport::Duration.parse. So in theory, we
+  # really only need to do the interval casting, but we're doing the
+  # SET LOCAL to be absolutely sure intervalstyle is correct.
+  #
+  # https://github.com/rails/rails/blob/v8.0.3/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb#L979-L980
+  def partition_interval_iso_8601
+    transaction do
+      connection.execute("SET LOCAL intervalstyle TO 'iso_8601'")
+      connection.select_value("SELECT #{connection.quote(partition_interval)}::interval")
+    end
+  end
+
+  def partition_rename_adapter
+    unless SUPPORTED_PARTITION_TYPES.include?(partition_type)
+      raise PgHaMigrations::InvalidPartmanConfigError,
+        "Expected partition_type to be in #{SUPPORTED_PARTITION_TYPES.inspect} " \
+        "but received #{partition_type.inspect}"
+    end
+
+    duration = ActiveSupport::Duration.parse(partition_interval_iso_8601)
+
+    if duration.parts.size != 1
+      raise PgHaMigrations::InvalidPartmanConfigError,
+        "Partition renaming for complex partition_interval #{duration.iso8601.inspect} not supported"
+    end
+
+    # Quarterly and weekly have special meaning in Partman 4 with
+    # specific datetime strings that need to be handled separately.
+    #
+    # The intervals "1 week" and "3 months" will not match the first
+    # two conditionals and will fallthrough to standard adapters below.
+    if duration == 1.week && datetime_string == "IYYY\"w\"IW"
+      PgHaMigrations::WeeklyPartmanRenameAdapter.new(self)
+    elsif duration == 3.months && datetime_string == "YYYY\"q\"Q"
+      PgHaMigrations::QuarterlyPartmanRenameAdapter.new(self)
+    elsif duration >= 1.year
+      PgHaMigrations::YearToForeverPartmanRenameAdapter.new(self)
+    elsif duration >= 1.month && duration < 1.year
+      PgHaMigrations::MonthToYearPartmanRenameAdapter.new(self)
+    elsif duration >= 1.day && duration < 1.month
+      PgHaMigrations::DayToMonthPartmanRenameAdapter.new(self)
+    elsif duration >= 1.minute && duration < 1.day
+      PgHaMigrations::MinuteToDayPartmanRenameAdapter.new(self)
+    elsif duration >= 1.second && duration < 1.minute
+      PgHaMigrations::SecondToMinutePartmanRenameAdapter.new(self)
+    else
+      raise PgHaMigrations::InvalidPartmanConfigError,
+        "Expected partition_interval to be greater than 1 second " \
+        "but received #{duration.iso8601.inspect}"
+    end
   end
 end

data/lib/pg_ha_migrations/partman_rename_adapter.rb

@@ -0,0 +1,209 @@
+module PgHaMigrations
+  class AbstractPartmanRenameAdapter
+    def initialize(part_config)
+      if part_config.datetime_string != source_datetime_string
+        raise PgHaMigrations::InvalidPartmanConfigError,
+          "Expected datetime_string to be #{source_datetime_string.inspect} " \
+          "but received #{part_config.datetime_string.inspect}"
+      end
+    end
+
+    def alter_table_sql(partitions)
+      sql = partitions.filter_map do |partition|
+        next if partition.name =~ /\A.+_default\z/
+
+        if partition.name !~ source_name_suffix_pattern
+          raise PgHaMigrations::InvalidIdentifierError,
+            "Expected #{partition.name.inspect} to match #{source_name_suffix_pattern.inspect}"
+        end
+
+        begin
+          "ALTER TABLE #{partition.fully_qualified_name} RENAME TO #{target_table_name(partition.name)};"
+        rescue Date::Error
+          raise PgHaMigrations::InvalidIdentifierError,
+            "Expected #{partition.name.inspect} suffix to be a parseable DateTime"
+        end
+      end.join("\n")
+
+      # This wraps the SQL in an anonymous function such that
+      # the statement timeout would apply to the entire batch of
+      # statements instead of each individual statement
+      "DO $$ BEGIN #{sql} END; $$;"
+    end
+
+    def target_table_name(table_name)
+      raise "#{__method__} should be implemented in subclass"
+    end
+
+    def source_datetime_string
+      raise "#{__method__} should be implemented in subclass"
+    end
+
+    def source_name_suffix_pattern
+      raise "#{__method__} should be implemented in subclass"
+    end
+
+    def target_datetime_string
+      raise "#{__method__} should be implemented in subclass"
+    end
+  end
+
+  class YearToForeverPartmanRenameAdapter < AbstractPartmanRenameAdapter
+    def target_table_name(table_name)
+      table_name + "0101"
+    end
+
+    def source_datetime_string
+      "YYYY"
+    end
+
+    def source_name_suffix_pattern
+      /\A.+_p\d{4}\z/
+    end
+
+    def target_datetime_string
+      "YYYYMMDD"
+    end
+  end
+
+   class QuarterlyPartmanRenameAdapter < AbstractPartmanRenameAdapter
+    QUARTER_MONTH_MAPPING = {
+      "1" => "01",
+      "2" => "04",
+      "3" => "07",
+      "4" => "10",
+    }
+
+    def target_table_name(table_name)
+      base_name = table_name[0...-6]
+
+      year = table_name.last(6).first(4)
+
+      month = QUARTER_MONTH_MAPPING.fetch(table_name.last(1))
+
+      base_name + year + month + "01"
+    end
+
+    def source_datetime_string
+      "YYYY\"q\"Q"
+    end
+
+    def source_name_suffix_pattern
+      /\A.+_p\d{4}q(1|2|3|4)\z/
+    end
+
+    def target_datetime_string
+      "YYYYMMDD"
+    end
+  end
+
+   class MonthToYearPartmanRenameAdapter < AbstractPartmanRenameAdapter
+    def target_table_name(table_name)
+      base_name = table_name[0...-7]
+
+      partition_datetime = DateTime.strptime(table_name.last(7), "%Y_%m")
+
+      base_name + partition_datetime.strftime("%Y%m%d")
+    end
+
+    def source_datetime_string
+      "YYYY_MM"
+    end
+
+    def source_name_suffix_pattern
+      /\A.+_p\d{4}_\d{2}\z/
+    end
+
+    def target_datetime_string
+      "YYYYMMDD"
+    end
+  end
+
+   class WeeklyPartmanRenameAdapter < AbstractPartmanRenameAdapter
+    def target_table_name(table_name)
+      base_name = table_name[0...-7]
+
+      partition_datetime = DateTime.strptime(table_name.last(7), "%Gw%V")
+
+      base_name + partition_datetime.strftime("%Y%m%d")
+    end
+
+    def source_datetime_string
+      "IYYY\"w\"IW"
+    end
+
+    def source_name_suffix_pattern
+      /\A.+_p\d{4}w\d{2}\z/
+    end
+
+    def target_datetime_string
+      "YYYYMMDD"
+    end
+  end
+
+   class DayToMonthPartmanRenameAdapter < AbstractPartmanRenameAdapter
+    def target_table_name(table_name)
+      base_name = table_name[0...-10]
+
+      partition_datetime = DateTime.strptime(table_name.last(10), "%Y_%m_%d")
+
+      base_name + partition_datetime.strftime("%Y%m%d")
+    end
+
+    def source_datetime_string
+      "YYYY_MM_DD"
+    end
+
+    def source_name_suffix_pattern
+      /\A.+_p\d{4}_\d{2}_\d{2}\z/
+    end
+
+    def target_datetime_string
+      "YYYYMMDD"
+    end
+  end
+
+   class MinuteToDayPartmanRenameAdapter < AbstractPartmanRenameAdapter
+    def target_table_name(table_name)
+      base_name = table_name[0...-15]
+
+      partition_datetime = DateTime.strptime(table_name.last(15), "%Y_%m_%d_%H%M")
+
+      base_name + partition_datetime.strftime("%Y%m%d_%H%M%S")
+    end
+
+    def source_datetime_string
+      "YYYY_MM_DD_HH24MI"
+    end
+
+    def source_name_suffix_pattern
+      /\A.+_p\d{4}_\d{2}_\d{2}_\d{4}\z/
+    end
+
+    def target_datetime_string
+      "YYYYMMDD_HH24MISS"
+    end
+  end
+
+   class SecondToMinutePartmanRenameAdapter < AbstractPartmanRenameAdapter
+    def target_table_name(table_name)
+      base_name = table_name[0...-17]
+
+      partition_datetime = DateTime.strptime(table_name.last(17), "%Y_%m_%d_%H%M%S")
+
+      base_name + partition_datetime.strftime("%Y%m%d_%H%M%S")
+    end
+
+    def source_datetime_string
+      "YYYY_MM_DD_HH24MISS"
+    end
+
+    def source_name_suffix_pattern
+      /\A.+_p\d{4}_\d{2}_\d{2}_\d{6}\z/
+    end
+
+    def target_datetime_string
+      "YYYYMMDD_HH24MISS"
+    end
+  end
+end

data/lib/pg_ha_migrations/relation.rb

@@ -1,4 +1,8 @@
 module PgHaMigrations
+  # This object represents a pointer to an actual relation in Postgres.
+  # The mode attribute is optional metadata which can represent a lock
+  # that has already been acquired or a potential lock that we are
+  # looking to acquire.
   Relation = Struct.new(:name, :schema, :mode) do
     def self.connection
       ActiveRecord::Base.connection
@@ -13,12 +17,6 @@ module PgHaMigrations
       self.mode = LockMode.new(mode) if mode.present?
     end
 
-    def conflicts_with?(other)
-      self == other && (
-        mode.nil? || other.mode.nil? || mode.conflicts_with?(other.mode)
-      )
-    end
-
     def fully_qualified_name
       @fully_qualified_name ||= [
         PG::Connection.quote_ident(schema),
@@ -30,8 +28,26 @@ module PgHaMigrations
       name.present? && schema.present?
     end
 
+    def conflicts_with?(other)
+      eql?(other) && (
+        mode.nil? || other.mode.nil? || mode.conflicts_with?(other.mode)
+      )
+    end
+
+    # Some code paths need to compare lock modes, while others just need
+    # to determine if the relation is the same object in Postgres, so
+    # equality here is simply looking at the relation name / schema.
+    # To also compare lock modes, #conflicts_with? is used.
+    def eql?(other)
+      other.is_a?(Relation) && hash == other.hash
+    end
+
     def ==(other)
-      other.is_a?(Relation) && name == other.name && schema == other.schema
+      eql?(other)
+    end
+
+    def hash
+      [name, schema].hash
     end
   end
 
@@ -97,6 +113,18 @@ module PgHaMigrations
       tables
     end
 
+    def check_constraints
+      connection.structs_from_sql(PgHaMigrations::CheckConstraint, <<~SQL)
+        SELECT conname AS name, pg_get_constraintdef(pg_constraint.oid) AS definition, convalidated AS validated
+        FROM pg_constraint, pg_class, pg_namespace
+        WHERE pg_class.oid = pg_constraint.conrelid
+          AND pg_class.relnamespace = pg_namespace.oid
+          AND pg_class.relname = #{connection.quote(name)}
+          AND pg_namespace.nspname = #{connection.quote(schema)}
+          AND pg_constraint.contype = 'c' -- 'c' stands for check constraints
+      SQL
+    end
+
     def has_rows?
       connection.select_value("SELECT EXISTS (SELECT 1 FROM #{fully_qualified_name} LIMIT 1)")
     end
@@ -111,6 +139,30 @@ module PgHaMigrations
     end
   end
 
+  class PartmanTable < Table
+    IDENTIFIER_REGEX = /^[a-z_][a-z_\d]*$/
+
+    def initialize(name, schema, mode=nil)
+      if name !~ IDENTIFIER_REGEX
+        raise InvalidIdentifierError, "Partman requires table names to be lowercase with underscores"
+      end
+
+      if schema !~ IDENTIFIER_REGEX
+        raise InvalidIdentifierError, "Partman requires schema names to be lowercase with underscores"
+      end
+
+      super
+    end
+
+    def fully_qualified_name
+      "#{schema}.#{name}"
+    end
+
+    def part_config(partman_extension:)
+      PgHaMigrations::PartmanConfig.find(fully_qualified_name, partman_extension: partman_extension)
+    end
+  end
+
   class Index < Relation
     MAX_NAME_SIZE = 63 # bytes
 
@@ -152,4 +204,45 @@ module PgHaMigrations
       SQL
     end
   end
+
+  class TableCollection
+    include Enumerable
+
+    attr_reader :raw_set
+
+    delegate :each, to: :raw_set
+    delegate :mode, to: :first
+
+    def self.from_table_names(tables, mode=nil)
+      new(tables) { |table| Table.from_table_name(table, mode) }
+    end
+
+    def initialize(tables, &blk)
+      @raw_set = tables.map(&blk).to_set
+
+      if raw_set.empty?
+        raise ArgumentError, "Expected a non-empty list of tables"
+      end
+
+      if raw_set.uniq(&:mode).size > 1
+        raise ArgumentError, "Expected all tables in collection to have the same lock mode"
+      end
+    end
+
+    def subset?(other)
+      raw_set.subset?(other.raw_set)
+    end
+
+    def to_sql
+      map(&:fully_qualified_name).join(", ")
+    end
+
+    def with_partitions
+      tables = flat_map do |table|
+        table.partitions(include_sub_partitions: true, include_self: true)
+      end
+
+      self.class.new(tables)
+    end
+  end
 end