pg_ha_migrations 1.6.0 → 1.8.0

data/lib/pg_ha_migrations/lock_mode.rb

@@ -0,0 +1,100 @@
+module PgHaMigrations
+  class LockMode
+    include Comparable
+
+    MODE_CONFLICTS = ActiveSupport::OrderedHash.new
+
+    MODE_CONFLICTS[:access_share] = %i[
+      access_exclusive
+    ]
+
+    MODE_CONFLICTS[:row_share] = %i[
+      exclusive
+      access_exclusive
+    ]
+
+    MODE_CONFLICTS[:row_exclusive] = %i[
+      share
+      share_row_exclusive
+      exclusive
+      access_exclusive
+    ]
+
+    MODE_CONFLICTS[:share_update_exclusive] = %i[
+      share_update_exclusive
+      share
+      share_row_exclusive
+      exclusive
+      access_exclusive
+    ]
+
+    MODE_CONFLICTS[:share] = %i[
+      row_exclusive
+      share_update_exclusive
+      share_row_exclusive
+      exclusive
+      access_exclusive
+    ]
+
+    MODE_CONFLICTS[:share_row_exclusive] = %i[
+      row_exclusive
+      share_update_exclusive
+      share
+      share_row_exclusive
+      exclusive
+      access_exclusive
+    ]
+
+    MODE_CONFLICTS[:exclusive] = %i[
+      row_share
+      row_exclusive
+      share_update_exclusive
+      share
+      share_row_exclusive
+      exclusive
+      access_exclusive
+    ]
+
+    MODE_CONFLICTS[:access_exclusive] = %i[
+      access_share
+      row_share
+      row_exclusive
+      share_update_exclusive
+      share
+      share_row_exclusive
+      exclusive
+      access_exclusive
+    ]
+
+    attr_reader :mode
+
+    delegate :to_s, to: :mode
+
+    def initialize(mode)
+      @mode = mode
+        .to_s
+        .underscore
+        .delete_suffix("_lock")
+        .to_sym
+
+      if !MODE_CONFLICTS.keys.include?(@mode)
+        raise ArgumentError, "Unrecognized lock mode #{@mode.inspect}. Valid modes: #{MODE_CONFLICTS.keys}"
+      end
+    end
+
+    def to_sql
+      mode
+        .to_s
+        .upcase
+        .gsub("_", " ")
+    end
+
+    def <=>(other)
+      MODE_CONFLICTS.keys.index(mode) <=> MODE_CONFLICTS.keys.index(other.mode)
+    end
+
+    def conflicts_with?(other)
+      MODE_CONFLICTS[mode].include?(other.mode)
+    end
+  end
+end

data/lib/pg_ha_migrations/partman_config.rb

@@ -0,0 +1,11 @@
+# This is an internal class that is not meant to be used directly
+class PgHaMigrations::PartmanConfig < ActiveRecord::Base
+  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"
+  end
+end

data/lib/pg_ha_migrations/relation.rb

@@ -0,0 +1,155 @@
+module PgHaMigrations
+  Relation = Struct.new(:name, :schema, :mode) do
+    def self.connection
+      ActiveRecord::Base.connection
+    end
+
+    delegate :inspect, to: :name
+    delegate :connection, to: :class
+
+    def initialize(name, schema, mode=nil)
+      super(name, schema)
+
+      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),
+        PG::Connection.quote_ident(name),
+      ].join(".")
+    end
+
+    def present?
+      name.present? && schema.present?
+    end
+
+    def ==(other)
+      other.is_a?(Relation) && name == other.name && schema == other.schema
+    end
+  end
+
+  class Table < Relation
+    def self.from_table_name(table, mode=nil)
+      pg_name = ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(table.to_s)
+
+      schema_conditional = if pg_name.schema
+        "#{connection.quote(pg_name.schema)}"
+      else
+        "ANY (current_schemas(false))"
+      end
+
+      schema = connection.select_value(<<~SQL)
+        SELECT schemaname
+        FROM pg_tables
+        WHERE tablename = #{connection.quote(pg_name.identifier)} AND schemaname = #{schema_conditional}
+        ORDER BY array_position(current_schemas(false), schemaname)
+        LIMIT 1
+      SQL
+
+      raise UndefinedTableError, "Table #{pg_name.quoted} does not exist#{" in search path" unless pg_name.schema}" unless schema.present?
+
+      new(pg_name.identifier, schema, mode)
+    end
+
+    def natively_partitioned?
+      return @natively_partitioned if defined?(@natively_partitioned)
+
+      @natively_partitioned = !!connection.select_value(<<~SQL)
+        SELECT true
+        FROM pg_partitioned_table, pg_class, pg_namespace
+        WHERE pg_class.oid = pg_partitioned_table.partrelid
+          AND pg_class.relnamespace = pg_namespace.oid
+          AND pg_class.relname = #{connection.quote(name)}
+          AND pg_namespace.nspname = #{connection.quote(schema)}
+      SQL
+    end
+
+    def partitions(include_sub_partitions: false, include_self: false)
+      tables = connection.structs_from_sql(self.class, <<~SQL)
+        SELECT child.relname AS name, child_ns.nspname AS schema, NULLIF('#{mode}', '') AS mode
+        FROM pg_inherits
+          JOIN pg_class parent        ON pg_inherits.inhparent = parent.oid
+          JOIN pg_class child         ON pg_inherits.inhrelid  = child.oid
+          JOIN pg_namespace parent_ns ON parent.relnamespace = parent_ns.oid
+          JOIN pg_namespace child_ns  ON child.relnamespace = child_ns.oid
+        WHERE parent.relname = #{connection.quote(name)}
+          AND parent_ns.nspname = #{connection.quote(schema)}
+        ORDER BY child.oid -- Ensure consistent ordering for tests
+      SQL
+
+      if include_sub_partitions
+        sub_partitions = tables.each_with_object([]) do |table, arr|
+          arr.concat(table.partitions(include_sub_partitions: true))
+        end
+
+        tables.concat(sub_partitions)
+      end
+
+      tables.prepend(self) if include_self
+
+      tables
+    end
+
+    def has_rows?
+      connection.select_value("SELECT EXISTS (SELECT 1 FROM #{fully_qualified_name} LIMIT 1)")
+    end
+
+    def total_bytes
+      connection.select_value(<<~SQL)
+        SELECT pg_total_relation_size(pg_class.oid)
+        FROM pg_class, pg_namespace
+        WHERE pg_class.relname = #{connection.quote(name)}
+          AND pg_namespace.nspname = #{connection.quote(schema)}
+      SQL
+    end
+  end
+
+  class Index < Relation
+    MAX_NAME_SIZE = 63 # bytes
+
+    def self.from_table_and_columns(table, columns)
+      name = connection.index_name(table.name, columns)
+
+      # modified from https://github.com/rails/rails/pull/47753
+      if name.bytesize > MAX_NAME_SIZE
+        hashed_identifier = "_#{OpenSSL::Digest::SHA256.hexdigest(name).first(10)}"
+        description = name.sub("index_#{table.name}_on", "idx_on")
+
+        short_limit = MAX_NAME_SIZE - hashed_identifier.bytesize
+        short_description = description.mb_chars.limit(short_limit).to_s
+
+        name = "#{short_description}#{hashed_identifier}"
+      end
+
+      new(name, table)
+    end
+
+    attr_accessor :table
+
+    def initialize(name, table)
+      super(name, table.schema)
+
+      self.table = table
+
+      connection.send(:validate_index_length!, table.name, name)
+    end
+
+    def valid?
+      !!connection.select_value(<<~SQL)
+        SELECT pg_index.indisvalid
+        FROM pg_index, pg_class, pg_namespace
+        WHERE pg_class.oid = pg_index.indexrelid
+          AND pg_class.relnamespace = pg_namespace.oid
+          AND pg_namespace.nspname = #{connection.quote(schema)}
+          AND pg_class.relname = #{connection.quote(name)}
+      SQL
+    end
+  end
+end

data/lib/pg_ha_migrations/safe_statements.rb

@@ -144,6 +144,22 @@ module PgHaMigrations::SafeStatements
     end
   end
 
+  def safe_add_index_on_empty_table(table, columns, options={})
+    if options[:algorithm] == :concurrently
+      raise ArgumentError, "Cannot call safe_add_index_on_empty_table with :algorithm => :concurrently"
+    end
+
+    # Avoids taking out an unnecessary SHARE lock if the table does have data
+    ensure_small_table!(table, empty: true)
+
+    safely_acquire_lock_for_table(table, mode: :share) do
+      # Ensure data wasn't written in the split second after the first check
+      ensure_small_table!(table, empty: true)
+
+      unsafe_add_index(table, columns, **options)
+    end
+  end
+
   def safe_add_concurrent_index(table, columns, options={})
     unsafe_add_index(table, columns, **options.merge(:algorithm => :concurrently))
   end
@@ -160,26 +176,96 @@ module PgHaMigrations::SafeStatements
     unsafe_remove_index(table, **options.merge(:algorithm => :concurrently))
   end
 
-  def safe_set_maintenance_work_mem_gb(gigabytes)
-    unsafe_execute("SET maintenance_work_mem = '#{PG::Connection.escape_string(gigabytes.to_s)} GB'")
-  end
+  def safe_add_concurrent_partitioned_index(
+    table,
+    columns,
+    name: nil,
+    if_not_exists: nil,
+    using: nil,
+    unique: nil,
+    where: nil,
+    comment: nil
+  )
+
+    if ActiveRecord::Base.connection.postgresql_version < 11_00_00
+      raise PgHaMigrations::InvalidMigrationError, "Concurrent partitioned index creation not supported on Postgres databases before version 11"
+    end
 
-  def safe_add_unvalidated_check_constraint(table, expression, name:)
-    unsafe_add_check_constraint(table, expression, name: name, validate: false)
-  end
+    parent_table = PgHaMigrations::Table.from_table_name(table)
 
-  def unsafe_add_check_constraint(table, expression, name:, validate: true)
-    raise ArgumentError, "Expected <name> to be present" unless name.present?
+    raise PgHaMigrations::InvalidMigrationError, "Table #{parent_table.inspect} is not a partitioned table" unless parent_table.natively_partitioned?
 
-    quoted_table_name = connection.quote_table_name(table)
-    quoted_constraint_name = connection.quote_table_name(name)
-    sql = "ALTER TABLE #{quoted_table_name} ADD CONSTRAINT #{quoted_constraint_name} CHECK (#{expression}) #{validate ? "" : "NOT VALID"}"
+    parent_index = if name.present?
+      PgHaMigrations::Index.new(name, parent_table)
+    else
+      PgHaMigrations::Index.from_table_and_columns(parent_table, columns)
+    end
 
-    safely_acquire_lock_for_table(table) do
-      say_with_time "add_check_constraint(#{table.inspect}, #{expression.inspect}, name: #{name.inspect}, validate: #{validate.inspect})" do
-        connection.execute(sql)
+    # Short-circuit when if_not_exists: true and index already valid
+    return if if_not_exists && parent_index.valid?
+
+    child_indexes = parent_table.partitions.map do |child_table|
+      PgHaMigrations::Index.from_table_and_columns(child_table, columns)
+    end
+
+    # TODO: take out ShareLock after issue #39 is implemented
+    safely_acquire_lock_for_table(parent_table.fully_qualified_name) do
+      # CREATE INDEX ON ONLY parent_table
+      unsafe_add_index(
+        parent_table.fully_qualified_name,
+        columns,
+        name: parent_index.name,
+        if_not_exists: if_not_exists,
+        using: using,
+        unique: unique,
+        where: where,
+        comment: comment,
+        algorithm: :only, # see lib/pg_ha_migrations/hacks/add_index_on_only.rb
+      )
+    end
+
+    child_indexes.each do |child_index|
+      add_index_method = if child_index.table.natively_partitioned?
+        :safe_add_concurrent_partitioned_index
+      else
+        :safe_add_concurrent_index
+      end
+
+      send(
+        add_index_method,
+        child_index.table.fully_qualified_name,
+        columns,
+        name: child_index.name,
+        if_not_exists: if_not_exists,
+        using: using,
+        unique: unique,
+        where: where,
+      )
+    end
+
+    # Avoid taking out an unnecessary lock if there are no child tables to attach
+    if child_indexes.present?
+      safely_acquire_lock_for_table(parent_table.fully_qualified_name) do
+        child_indexes.each do |child_index|
+          say_with_time "Attaching index #{child_index.inspect} to #{parent_index.inspect}" do
+            connection.execute(<<~SQL)
+              ALTER INDEX #{parent_index.fully_qualified_name}
+              ATTACH PARTITION #{child_index.fully_qualified_name}
+            SQL
+          end
+        end
       end
     end
+
+    raise PgHaMigrations::InvalidMigrationError, "Unexpected state. Parent index #{parent_index.inspect} is invalid" unless parent_index.valid?
+  end
+
+  def safe_set_maintenance_work_mem_gb(gigabytes)
+    unsafe_execute("SET maintenance_work_mem = '#{PG::Connection.escape_string(gigabytes.to_s)} GB'")
+  end
+
+  def safe_add_unvalidated_check_constraint(table, expression, name:)
+    unsafe_add_check_constraint(table, expression, name: name, validate: false)
   end
 
   def safe_validate_check_constraint(table, name:)
@@ -224,6 +310,197 @@ module PgHaMigrations::SafeStatements
     end
   end
 
+  def safe_create_partitioned_table(table, partition_key:, type:, infer_primary_key: nil, **options, &block)
+    raise ArgumentError, "Expected <partition_key> to be present" unless partition_key.present?
+
+    unless PgHaMigrations::PARTITION_TYPES.include?(type)
+      raise ArgumentError, "Expected <type> to be symbol in #{PgHaMigrations::PARTITION_TYPES} but received #{type.inspect}"
+    end
+
+    if ActiveRecord::Base.connection.postgresql_version < 10_00_00
+      raise PgHaMigrations::InvalidMigrationError, "Native partitioning not supported on Postgres databases before version 10"
+    end
+
+    if type == :hash && ActiveRecord::Base.connection.postgresql_version < 11_00_00
+      raise PgHaMigrations::InvalidMigrationError, "Hash partitioning not supported on Postgres databases before version 11"
+    end
+
+    if infer_primary_key.nil?
+      infer_primary_key = PgHaMigrations.config.infer_primary_key_on_partitioned_tables
+    end
+
+    # Newer versions of Rails will set the primary key column to the type :primary_key.
+    # This performs some extra logic that we can't easily undo which causes problems when
+    # trying to inject the partition key into the PK. Now, it would be nice to lookup the
+    # default primary key type instead of simply using :bigserial, but it doesn't appear
+    # that we have access to the Rails configuration from within our migrations.
+    if options[:id].nil? || options[:id] == :primary_key
+      options[:id] = :bigserial
+    end
+
+    quoted_partition_key = if partition_key.is_a?(Proc)
+      # Lambda syntax, like in other migration methods, implies an expression that
+      # cannot be easily sanitized.
+      #
+      # e.g ->{ "(created_at::date)" }
+      partition_key.call.to_s
+    else
+      # Otherwise, assume key is a column name or array of column names
+      Array.wrap(partition_key).map { |col| connection.quote_column_name(col) }.join(",")
+    end
+
+    options[:options] = "PARTITION BY #{type.upcase} (#{quoted_partition_key})"
+
+    safe_create_table(table, options) do |td|
+      block.call(td) if block
+
+      next unless options[:id]
+
+      pk_columns = td.columns.each_with_object([]) do |col, arr|
+        next unless col.options[:primary_key]
+
+        col.options[:primary_key] = false
+
+        arr << col.name
+      end
+
+      if infer_primary_key && !partition_key.is_a?(Proc) && ActiveRecord::Base.connection.postgresql_version >= 11_00_00
+        td.primary_keys(pk_columns.concat(Array.wrap(partition_key)).map(&:to_s).uniq)
+      end
+    end
+  end
+
+  def safe_partman_create_parent(table, **options)
+    if options[:retention].present? || options[:retention_keep_table] == false
+      raise PgHaMigrations::UnsafeMigrationError.new(":retention and/or :retention_keep_table => false can potentially result in data loss if misconfigured. Please use unsafe_partman_create_parent if you want to set these options")
+    end
+
+    unsafe_partman_create_parent(table, **options)
+  end
+
+  def unsafe_partman_create_parent(
+    table,
+    partition_key:,
+    interval:,
+    infinite_time_partitions: true,
+    inherit_privileges: true,
+    premake: nil,
+    start_partition: nil,
+    template_table: nil,
+    retention: nil,
+    retention_keep_table: nil
+  )
+    raise ArgumentError, "Expected <partition_key> to be present" unless partition_key.present?
+    raise ArgumentError, "Expected <interval> to be present" unless interval.present?
+
+    if ActiveRecord::Base.connection.postgresql_version < 11_00_00
+      raise PgHaMigrations::InvalidMigrationError, "Native partitioning with partman not supported on Postgres databases before version 11"
+    end
+
+    formatted_start_partition = nil
+
+    if start_partition.present?
+      if !start_partition.is_a?(Date) && !start_partition.is_a?(Time) && !start_partition.is_a?(DateTime)
+        raise PgHaMigrations::InvalidMigrationError, "Expected <start_partition> to be Date, Time, or DateTime object but received #{start_partition.class}"
+      end
+
+      formatted_start_partition = if start_partition.respond_to?(:to_fs)
+        start_partition.to_fs(:db)
+      else
+        start_partition.to_s(:db)
+      end
+    end
+
+    create_parent_options = {
+      parent_table: _fully_qualified_table_name_for_partman(table),
+      template_table: template_table ? _fully_qualified_table_name_for_partman(template_table) : nil,
+      control: partition_key,
+      type: "native",
+      interval: interval,
+      premake: premake,
+      start_partition: formatted_start_partition,
+    }.compact
+
+    create_parent_sql = create_parent_options.map { |k, v| "p_#{k} := #{connection.quote(v)}" }.join(", ")
+
+    log_message = "partman_create_parent(#{table.inspect}, " \
+      "partition_key: #{partition_key.inspect}, " \
+      "interval: #{interval.inspect}, " \
+      "premake: #{premake.inspect}, " \
+      "start_partition: #{start_partition.inspect}, " \
+      "template_table: #{template_table.inspect})"
+
+    say_with_time(log_message) do
+      connection.execute("SELECT #{_quoted_partman_schema}.create_parent(#{create_parent_sql})")
+    end
+
+    update_config_options = {
+      infinite_time_partitions: infinite_time_partitions,
+      inherit_privileges: inherit_privileges,
+      retention: retention,
+      retention_keep_table: retention_keep_table,
+    }.compact
+
+    unsafe_partman_update_config(table, **update_config_options)
+  end
+
+  def safe_partman_update_config(table, **options)
+    if options[:retention].present? || options[:retention_keep_table] == false
+      raise PgHaMigrations::UnsafeMigrationError.new(":retention and/or :retention_keep_table => false can potentially result in data loss if misconfigured. Please use unsafe_partman_update_config if you want to set these options")
+    end
+
+    unsafe_partman_update_config(table, **options)
+  end
+
+  def unsafe_partman_update_config(table, **options)
+    invalid_options = options.keys - PgHaMigrations::PARTMAN_UPDATE_CONFIG_OPTIONS
+
+    raise ArgumentError, "Unrecognized argument(s): #{invalid_options}" unless invalid_options.empty?
+
+    PgHaMigrations::PartmanConfig.schema = _quoted_partman_schema
+
+    config = PgHaMigrations::PartmanConfig.find(_fully_qualified_table_name_for_partman(table))
+
+    config.assign_attributes(**options)
+
+    inherit_privileges_changed = config.inherit_privileges_changed?
+
+    say_with_time "partman_update_config(#{table.inspect}, #{options.map { |k,v| "#{k}: #{v.inspect}" }.join(", ")})" do
+      config.save!
+    end
+
+    safe_partman_reapply_privileges(table) if inherit_privileges_changed
+  end
+
+  def safe_partman_reapply_privileges(table)
+    say_with_time "partman_reapply_privileges(#{table.inspect})" do
+      connection.execute("SELECT #{_quoted_partman_schema}.reapply_privileges('#{_fully_qualified_table_name_for_partman(table)}')")
+    end
+  end
+
+  def _quoted_partman_schema
+    schema = connection.select_value(<<~SQL)
+      SELECT nspname
+      FROM pg_namespace JOIN pg_extension
+        ON pg_namespace.oid = pg_extension.extnamespace
+      WHERE pg_extension.extname = 'pg_partman'
+    SQL
+
+    raise PgHaMigrations::InvalidMigrationError, "The pg_partman extension is not installed" unless schema.present?
+
+    connection.quote_schema_name(schema)
+  end
+
+  def _fully_qualified_table_name_for_partman(table)
+    table = PgHaMigrations::Table.from_table_name(table)
+
+    [table.schema, table.name].each do |identifier|
+      if identifier.to_s !~ /^[a-z_][a-z_\d]*$/
+        raise PgHaMigrations::InvalidMigrationError, "Partman requires schema / table names to be lowercase with underscores"
+      end
+    end.join(".")
+  end
+
   def _per_migration_caller
     @_per_migration_caller ||= Kernel.caller
   end
@@ -251,17 +528,39 @@ module PgHaMigrations::SafeStatements
     super(conn, direction)
   end
 
-  def safely_acquire_lock_for_table(table, &block)
+  def safely_acquire_lock_for_table(table, mode: :access_exclusive, &block)
+    nested_target_table = Thread.current[__method__]
+
     _check_postgres_adapter!
-    table = table.to_s
-    quoted_table_name = connection.quote_table_name(table)
+
+    target_table = PgHaMigrations::Table.from_table_name(table, mode)
+
+    if nested_target_table
+      if nested_target_table != target_table
+        raise PgHaMigrations::InvalidMigrationError, "Nested lock detected! Cannot acquire lock on #{target_table.fully_qualified_name} while #{nested_target_table.fully_qualified_name} is locked."
+      elsif nested_target_table.mode < target_table.mode
+        raise PgHaMigrations::InvalidMigrationError, "Lock escalation detected! Cannot change lock level from :#{nested_target_table.mode} to :#{target_table.mode} for #{target_table.fully_qualified_name}."
+      end
+    else
+      Thread.current[__method__] = target_table
+    end
+
+    # Locking a partitioned table will also lock child tables (including sub-partitions),
+    # so we need to check for blocking queries on those tables as well
+    target_tables = target_table.partitions(include_sub_partitions: true, include_self: true)
 
     successfully_acquired_lock = false
 
     until successfully_acquired_lock
       while (
         blocking_transactions = PgHaMigrations::BlockingDatabaseTransactions.find_blocking_transactions("#{PgHaMigrations::LOCK_TIMEOUT_SECONDS} seconds")
-        blocking_transactions.any? { |query| query.tables_with_locks.include?(table) }
+        blocking_transactions.any? do |query|
+          query.tables_with_locks.any? do |locked_table|
+            target_tables.any? do |target_table|
+              target_table.conflicts_with?(locked_table)
+            end
+          end
+        end
       )
         say "Waiting on blocking transactions:"
         blocking_transactions.each do |blocking_transaction|
@@ -274,13 +573,13 @@ module PgHaMigrations::SafeStatements
         adjust_timeout_method = connection.postgresql_version >= 9_03_00 ? :adjust_lock_timeout : :adjust_statement_timeout
         begin
           method(adjust_timeout_method).call(PgHaMigrations::LOCK_TIMEOUT_SECONDS) do
-            connection.execute("LOCK #{quoted_table_name};")
+            connection.execute("LOCK #{target_table.fully_qualified_name} IN #{target_table.mode.to_sql} MODE;")
           end
           successfully_acquired_lock = true
         rescue ActiveRecord::StatementInvalid => e
           if e.message =~ /PG::LockNotAvailable.+ lock timeout/ || e.message =~ /PG::QueryCanceled.+ statement timeout/
             sleep_seconds = PgHaMigrations::LOCK_FAILURE_RETRY_DELAY_MULTLIPLIER * PgHaMigrations::LOCK_TIMEOUT_SECONDS
-            say "Timed out trying to acquire an exclusive lock on the #{quoted_table_name} table."
+            say "Timed out trying to acquire #{target_table.mode.to_sql} lock on the #{target_table.fully_qualified_name} table."
             say "Sleeping for #{sleep_seconds}s to allow potentially queued up queries to finish before continuing."
             sleep(sleep_seconds)
 
@@ -295,6 +594,8 @@ module PgHaMigrations::SafeStatements
         end
       end
     end
+  ensure
+    Thread.current[__method__] = nil unless nested_target_table
   end
 
   def adjust_lock_timeout(timeout_seconds = PgHaMigrations::LOCK_TIMEOUT_SECONDS, &block)
@@ -336,4 +637,16 @@ module PgHaMigrations::SafeStatements
       end
     end
   end
+
+  def ensure_small_table!(table, empty: false, threshold: PgHaMigrations::SMALL_TABLE_THRESHOLD_BYTES)
+    table = PgHaMigrations::Table.from_table_name(table)
+
+    if empty && table.has_rows?
+      raise PgHaMigrations::InvalidMigrationError, "Table #{table.inspect} has rows"
+    end
+
+    if table.total_bytes > threshold
+      raise PgHaMigrations::InvalidMigrationError, "Table #{table.inspect} is larger than #{threshold} bytes"
+    end
+  end
 end