pg_ha_migrations 1.7.0 → 2.0.0

data/lib/pg_ha_migrations/safe_statements.rb

@@ -1,24 +1,14 @@
 module PgHaMigrations::SafeStatements
-  PARTITION_TYPES = %i[range list hash]
-
-  PARTMAN_UPDATE_CONFIG_OPTIONS = %i[
-    infinite_time_partitions
-    inherit_privileges
-    premake
-    retention
-    retention_keep_table
-  ]
-
   def safe_added_columns_without_default_value
     @safe_added_columns_without_default_value ||= []
   end
 
-  def safe_create_table(table, options={}, &block)
+  def safe_create_table(table, **options, &block)
     if options[:force]
       raise PgHaMigrations::UnsafeMigrationError.new(":force is NOT SAFE! Explicitly call unsafe_drop_table first if you want to recreate an existing table")
     end
 
-    unsafe_create_table(table, options, &block)
+    unsafe_create_table(table, **options, &block)
   end
 
   def safe_create_enum_type(name, values=nil)
@@ -26,28 +16,20 @@ module PgHaMigrations::SafeStatements
     when nil
       raise ArgumentError, "safe_create_enum_type expects a set of values; if you want an enum with no values please pass an empty array"
     when []
-      unsafe_execute("CREATE TYPE #{PG::Connection.quote_ident(name.to_s)} AS ENUM ()")
+      raw_execute("CREATE TYPE #{PG::Connection.quote_ident(name.to_s)} AS ENUM ()")
     else
       escaped_values = values.map do |value|
         "'#{PG::Connection.escape_string(value.to_s)}'"
       end
-      unsafe_execute("CREATE TYPE #{PG::Connection.quote_ident(name.to_s)} AS ENUM (#{escaped_values.join(',')})")
+      raw_execute("CREATE TYPE #{PG::Connection.quote_ident(name.to_s)} AS ENUM (#{escaped_values.join(',')})")
     end
   end
 
   def safe_add_enum_value(name, value)
-    unsafe_execute("ALTER TYPE #{PG::Connection.quote_ident(name.to_s)} ADD VALUE '#{PG::Connection.escape_string(value)}'")
+    raw_execute("ALTER TYPE #{PG::Connection.quote_ident(name.to_s)} ADD VALUE '#{PG::Connection.escape_string(value)}'")
   end
 
-  def unsafe_rename_enum_value(name, old_value, new_value)
-    if ActiveRecord::Base.connection.postgresql_version < 10_00_00
-      raise PgHaMigrations::InvalidMigrationError, "Renaming an enum value is not supported on Postgres databases before version 10"
-    end
-
-    unsafe_execute("ALTER TYPE #{PG::Connection.quote_ident(name.to_s)} RENAME VALUE '#{PG::Connection.escape_string(old_value)}' TO '#{PG::Connection.escape_string(new_value)}'")
-  end
-
-  def safe_add_column(table, column, type, options = {})
+  def safe_add_column(table, column, type, **options)
     # Note: we don't believe we need to consider the odd case where
     # `:default => nil` or `:default => -> { null }` (or similar) is
     # passed because:
@@ -63,20 +45,14 @@ module PgHaMigrations::SafeStatements
         raise PgHaMigrations::UnsafeMigrationError.new(":default is not safe if the default value is volatile. Use safe_change_column_default afterwards then backfill the data to prevent locking the table")
       end
     elsif options[:null] == false
-      raise PgHaMigrations::UnsafeMigrationError.new(":null => false is NOT SAFE if the table has data! If you _really_ want to do this, use unsafe_make_column_not_nullable")
+      raise PgHaMigrations::UnsafeMigrationError.new(":null => false is NOT SAFE if the table has data! If you want to do this, use safe_make_column_not_nullable")
     end
 
     unless options.has_key?(:default)
       self.safe_added_columns_without_default_value << [table.to_s, column.to_s]
     end
 
-    unsafe_add_column(table, column, type, options)
-  end
-
-  def unsafe_add_column(table, column, type, options = {})
-    safely_acquire_lock_for_table(table) do
-      super(table, column, type, **options)
-    end
+    unsafe_add_column(table, column, type, **options)
   end
 
   def safe_change_column_default(table_name, column_name, default_value)
@@ -143,22 +119,131 @@ module PgHaMigrations::SafeStatements
   end
 
   def safe_make_column_nullable(table, column)
+    quoted_table_name = connection.quote_table_name(table)
+    quoted_column_name = connection.quote_column_name(column)
+
     safely_acquire_lock_for_table(table) do
-      unsafe_execute "ALTER TABLE #{table} ALTER COLUMN #{column} DROP NOT NULL"
+      raw_execute "ALTER TABLE #{quoted_table_name} ALTER COLUMN #{quoted_column_name} DROP NOT NULL"
     end
   end
 
-  def unsafe_make_column_not_nullable(table, column, options={}) # options arg is only present for backwards compatiblity
+  # Postgres 12+ can use a valid CHECK constraint to validate that no values of a column are null, avoiding
+  # a full table scan while holding an exclusive lock on the table when altering a column to NOT NULL
+  #
+  # Source:
+  # https://dba.stackexchange.com/questions/267947/how-can-i-set-a-column-to-not-null-without-locking-the-table-during-a-table-scan/268128#268128
+  # (https://archive.is/X55up)
+  def safe_make_column_not_nullable(table, column)
+    if ActiveRecord::Base.connection.postgresql_version < 12_00_00
+      raise PgHaMigrations::InvalidMigrationError, "Cannot safely make a column non-nullable before Postgres 12"
+    end
+
+    validated_table = PgHaMigrations::Table.from_table_name(table)
+    tmp_constraint_name = "tmp_not_null_constraint_#{OpenSSL::Digest::SHA256.hexdigest(column.to_s).first(7)}"
+
+    if validated_table.check_constraints.any? { |c| c.name == tmp_constraint_name }
+      raise PgHaMigrations::InvalidMigrationError, "A constraint #{tmp_constraint_name.inspect} already exists. " \
+        "This implies that a previous invocation of this method failed and left behind a temporary constraint. " \
+        "Please drop the constraint before attempting to run this method again."
+    end
+
+    safe_add_unvalidated_check_constraint(table, "#{connection.quote_column_name(column)} IS NOT NULL", name: tmp_constraint_name)
+    safe_validate_check_constraint(table, name: tmp_constraint_name)
+
+    # "Ordinarily this is checked during the ALTER TABLE by scanning the entire table; however, if a
+    # valid CHECK constraint is found which proves no NULL can exist, then the table scan is
+    # skipped."
+    # See: https://www.postgresql.org/docs/current/sql-altertable.html#SQL-ALTERTABLE-DESC-SET-DROP-NOT-NULL
+    unsafe_make_column_not_nullable(table, column)
+    unsafe_remove_constraint(table, name: tmp_constraint_name)
+  end
+
+  # This method is a variant of `safe_make_column_not_nullable` that is expected to always be fast;
+  # i.e., it will not perform a full table scan to check for null values.
+  def safe_make_column_not_nullable_from_check_constraint(table, column, constraint_name:, drop_constraint: true)
+    unless ActiveRecord::Base.connection.postgresql_version >= 12_00_00
+      raise PgHaMigrations::InvalidMigrationError, "Cannot safely make a column non-nullable before Postgres 12"
+    end
+
+    unless constraint_name
+      raise ArgumentError, "Expected <constraint_name> to be present"
+    end
+    constraint_name = constraint_name.to_s
+
+    quoted_table_name = connection.quote_table_name(table)
+    quoted_column_name = connection.quote_column_name(column)
+
+    validated_table = PgHaMigrations::Table.from_table_name(table)
+    constraint = validated_table.check_constraints.find do |c|
+      c.name == constraint_name
+    end
+
+    unless constraint
+      raise PgHaMigrations::InvalidMigrationError, "The provided constraint does not exist"
+    end
+
+    unless constraint.validated
+      raise PgHaMigrations::InvalidMigrationError, "The provided constraint is not validated"
+    end
+
+    # The constraint has to actually prove that no null values exist, so the
+    # constraint condition can't simply include the `IS NOT NULL` check. We
+    # don't try to handle all possible cases here. For example,
+    # `a IS NOT NULL AND b IS NOT NULL` would prove what we need, but it would
+    # be complicated to check. We must ensure, however, that we're not too
+    # loose. For example, `a IS NOT NULL OR b IS NOT NULL` would not prove that
+    # `a IS NOT NULL`.
+    unless constraint.definition =~ /\ACHECK \(*(#{Regexp.escape(column.to_s)}|#{Regexp.escape(quoted_column_name)}) IS NOT NULL\)*\Z/i
+      raise PgHaMigrations::InvalidMigrationError, "The provided constraint does not enforce non-null values for the column"
+    end
+
+    # We don't want to acquire an exclusive lock on the table twice, and we also don't want it to be
+    # posssible to have the NOT NULL constraint addition succeed while the constraint removal fails,
+    # so we acquire the lock once and do both operations in the same block.
     safely_acquire_lock_for_table(table) do
-      unsafe_execute "ALTER TABLE #{table} ALTER COLUMN #{column} SET NOT NULL"
+      # "Ordinarily this is checked during the ALTER TABLE by scanning the entire table; however, if a
+      # valid CHECK constraint is found which proves no NULL can exist, then the table scan is
+      # skipped."
+      # See: https://www.postgresql.org/docs/current/sql-altertable.html#SQL-ALTERTABLE-DESC-SET-DROP-NOT-NULL
+      unsafe_make_column_not_nullable(table, column)
+
+      if drop_constraint
+        unsafe_remove_constraint(table, name: constraint_name)
+      end
+    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
+
+    # Check if nulls_not_distinct was provided but PostgreSQL version doesn't support it
+    if options[:nulls_not_distinct] && ActiveRecord::Base.connection.postgresql_version < 15_00_00
+      raise PgHaMigrations::InvalidMigrationError, "nulls_not_distinct option requires PostgreSQL 15 or higher"
+    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={})
+  def safe_add_concurrent_index(table, columns, **options)
+    # Check if nulls_not_distinct was provided but PostgreSQL version doesn't support it
+    if options[:nulls_not_distinct] && ActiveRecord::Base.connection.postgresql_version < 15_00_00
+      raise PgHaMigrations::InvalidMigrationError, "nulls_not_distinct option requires PostgreSQL 15 or higher"
+    end
+
     unsafe_add_index(table, columns, **options.merge(:algorithm => :concurrently))
   end
 
-  def safe_remove_concurrent_index(table, options={})
+  def safe_remove_concurrent_index(table, **options)
     unless options.is_a?(Hash) && options.key?(:name)
       raise ArgumentError, "Expected safe_remove_concurrent_index to be called with arguments (table_name, :name => ...)"
     end
@@ -170,8 +255,96 @@ module PgHaMigrations::SafeStatements
     unsafe_remove_index(table, **options.merge(:algorithm => :concurrently))
   end
 
+  def safe_add_concurrent_partitioned_index(
+    table,
+    columns,
+    name: nil,
+    if_not_exists: nil,
+    using: nil,
+    unique: nil,
+    where: nil,
+    comment: nil,
+    nulls_not_distinct: nil
+  )
+    # Check if nulls_not_distinct was provided but PostgreSQL version doesn't support it
+    if !nulls_not_distinct.nil? && ActiveRecord::Base.connection.postgresql_version < 15_00_00
+      raise PgHaMigrations::InvalidMigrationError, "nulls_not_distinct option requires PostgreSQL 15 or higher"
+    end
+
+    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
+
+    parent_table = PgHaMigrations::Table.from_table_name(table)
+
+    raise PgHaMigrations::InvalidMigrationError, "Table #{parent_table.inspect} is not a partitioned table" unless parent_table.natively_partitioned?
+
+    parent_index = if name.present?
+      PgHaMigrations::Index.new(name, parent_table)
+    else
+      PgHaMigrations::Index.from_table_and_columns(parent_table, columns)
+    end
+
+    # 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
+
+    # 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,
+      nulls_not_distinct: nulls_not_distinct,
+      where: where,
+      comment: comment,
+      algorithm: :only, # see lib/pg_ha_migrations/hacks/add_index_on_only.rb
+    )
+
+    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,
+        nulls_not_distinct: nulls_not_distinct,
+        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'")
+    raw_execute("SET maintenance_work_mem = '#{PG::Connection.escape_string(gigabytes.to_s)} GB'")
   end
 
   def safe_add_unvalidated_check_constraint(table, expression, name:)
@@ -206,25 +379,11 @@ module PgHaMigrations::SafeStatements
     end
   end
 
-  def unsafe_remove_constraint(table, name:)
-    raise ArgumentError, "Expected <name> to be present" unless name.present?
-
-    quoted_table_name = connection.quote_table_name(table)
-    quoted_constraint_name = connection.quote_table_name(name)
-    sql = "ALTER TABLE #{quoted_table_name} DROP CONSTRAINT #{quoted_constraint_name}"
-
-    safely_acquire_lock_for_table(table) do
-      say_with_time "remove_constraint(#{table.inspect}, name: #{name.inspect})" do
-        connection.execute(sql)
-      end
-    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 PARTITION_TYPES.include?(type)
-      raise ArgumentError, "Expected <type> to be symbol in #{PARTITION_TYPES} but received #{type.inspect}"
+    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
@@ -261,7 +420,7 @@ module PgHaMigrations::SafeStatements
 
     options[:options] = "PARTITION BY #{type.upcase} (#{quoted_partition_key})"
 
-    safe_create_table(table, options) do |td|
+    safe_create_table(table, **options) do |td|
       block.call(td) if block
 
       next unless options[:id]
@@ -280,15 +439,7 @@ module PgHaMigrations::SafeStatements
     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(
+  def safe_partman_create_parent(
     table,
     partition_key:,
     interval:,
@@ -351,7 +502,7 @@ module PgHaMigrations::SafeStatements
       retention_keep_table: retention_keep_table,
     }.compact
 
-    unsafe_partman_update_config(create_parent_options[:parent_table], **update_config_options)
+    unsafe_partman_update_config(table, **update_config_options)
   end
 
   def safe_partman_update_config(table, **options)
@@ -362,26 +513,6 @@ module PgHaMigrations::SafeStatements
     unsafe_partman_update_config(table, **options)
   end
 
-  def unsafe_partman_update_config(table, **options)
-    invalid_options = options.keys - 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)}')")
@@ -402,38 +533,13 @@ module PgHaMigrations::SafeStatements
   end
 
   def _fully_qualified_table_name_for_partman(table)
-    identifiers = table.to_s.split(".")
-
-    raise PgHaMigrations::InvalidMigrationError, "Expected table to be in the format <table> or <schema>.<table> but received #{table}" if identifiers.size > 2
-
-    identifiers.each { |identifier| _validate_partman_identifier(identifier) }
-
-    schema_conditional = if identifiers.size > 1
-      "'#{identifiers.first}'"
-    else
-      "ANY (current_schemas(false))"
-    end
-
-    schema = connection.select_value(<<~SQL)
-      SELECT schemaname
-      FROM pg_tables
-      WHERE tablename = '#{identifiers.last}' AND schemaname = #{schema_conditional}
-      ORDER BY array_position(current_schemas(false), schemaname)
-      LIMIT 1
-    SQL
-
-    raise PgHaMigrations::InvalidMigrationError, "Could not find table #{table}" unless schema.present?
-
-    _validate_partman_identifier(schema)
-
-    # Quoting is unneeded since _validate_partman_identifier ensures the schema / table use standard naming conventions
-    "#{schema}.#{identifiers.last}"
-  end
+    table = PgHaMigrations::Table.from_table_name(table)
 
-  def _validate_partman_identifier(identifier)
-    if identifier.to_s !~ /^[a-z_][a-z_\d]*$/
-      raise PgHaMigrations::InvalidMigrationError, "Partman requires schema / table names to be lowercase with underscores"
-    end
+    [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
@@ -463,18 +569,50 @@ module PgHaMigrations::SafeStatements
     super(conn, direction)
   end
 
-  def safely_acquire_lock_for_table(table, &block)
+  def safely_acquire_lock_for_table(*tables, mode: :access_exclusive, &block)
     _check_postgres_adapter!
-    table = table.to_s
-    quoted_table_name = connection.quote_table_name(table)
+
+    target_tables = PgHaMigrations::TableCollection.from_table_names(tables, mode)
+
+    if @parent_lock_tables
+      if !target_tables.subset?(@parent_lock_tables)
+        raise PgHaMigrations::InvalidMigrationError,
+          "Nested lock detected! Cannot acquire lock on #{target_tables.to_sql} " \
+          "while #{@parent_lock_tables.to_sql} is locked."
+      end
+
+      if @parent_lock_tables.mode < target_tables.mode
+        raise PgHaMigrations::InvalidMigrationError,
+          "Lock escalation detected! Cannot change lock level from :#{@parent_lock_tables.mode} " \
+          "to :#{target_tables.mode} for #{target_tables.to_sql}."
+      end
+
+      # If in a nested context and all of the above checks have passed,
+      # we have already acquired the lock (or a lock at a higher level),
+      # and can simply execute the block and short-circuit.
+      block.call
+
+      return
+    end
 
     successfully_acquired_lock = false
 
     until successfully_acquired_lock
-      while (
+      loop do
         blocking_transactions = PgHaMigrations::BlockingDatabaseTransactions.find_blocking_transactions("#{PgHaMigrations::LOCK_TIMEOUT_SECONDS} seconds")
-        blocking_transactions.any? { |query| query.tables_with_locks.include?(table) }
-      )
+
+        # 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_with_partitions = target_tables.with_partitions
+
+        break unless blocking_transactions.any? do |query|
+          query.tables_with_locks.any? do |locked_table|
+            target_tables_with_partitions.any? do |target_table|
+              target_table.conflicts_with?(locked_table)
+            end
+          end
+        end
+
         say "Waiting on blocking transactions:"
         blocking_transactions.each do |blocking_transaction|
           say blocking_transaction.description
@@ -483,16 +621,21 @@ module PgHaMigrations::SafeStatements
       end
 
       connection.transaction do
-        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};")
+          # A lock timeout would apply to each individual table in the query,
+          # so we made a conscious decision to use a statement timeout here
+          # to keep behavior consistent in a multi-table lock scenario.
+          adjust_statement_timeout(PgHaMigrations::LOCK_TIMEOUT_SECONDS) do
+            connection.execute("LOCK #{target_tables.to_sql} IN #{target_tables.mode.to_sql} MODE;")
           end
           successfully_acquired_lock = true
         rescue ActiveRecord::StatementInvalid => e
+          # It is still possible to hit a lock timeout if the session has
+          # that value set to something less than LOCK_TIMEOUT_SECONDS.
+          # We should retry when either of these exceptions are raised.
           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_tables.mode.to_sql} lock on #{target_tables.to_sql}."
             say "Sleeping for #{sleep_seconds}s to allow potentially queued up queries to finish before continuing."
             sleep(sleep_seconds)
 
@@ -503,7 +646,13 @@ module PgHaMigrations::SafeStatements
         end
 
         if successfully_acquired_lock
-          block.call
+          @parent_lock_tables = target_tables
+
+          begin
+            block.call
+          ensure
+            @parent_lock_tables = nil
+          end
         end
       end
     end
@@ -548,4 +697,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