Method to safely add index on empty table

README.md

@@ -171,6 +171,14 @@ Unsafely make a column not nullable.
 unsafe_make_column_not_nullable :table, :column
 ```
 
+#### safe\_add\_index\_on\_empty\_table
+
+Safely add an index on a table with zero rows. This will raise an error if the table contains data.
+
+```ruby
+safe_add_index_on_empty_table :table, :column
+```
+
 #### safe\_add\_concurrent\_index
 
 Add an index concurrently.
@@ -476,6 +484,22 @@ Set maintenance work mem.
 safe_set_maintenance_work_mem_gb 1
 ```
 
+#### ensure\_small\_table!
+
+Ensure a table on disk is below the default threshold (10 megabytes).
+This will raise an error if the table is too large.
+
+```ruby
+ensure_small_table! :table
+```
+
+Ensure a table on disk is below a custom threshold and is empty.
+This will raise an error if the table is too large and/or contains data.
+
+```ruby
+ensure_small_table! :table, empty: true, threshold: 100.megabytes
+```
+
 ### Configuration
 
 The gem can be configured in an initializer.

lib/pg_ha_migrations.rb

@@ -3,6 +3,7 @@
 require "active_record"
 require "active_record/migration"
 require "active_record/connection_adapters/postgresql/utils"
+require "active_support/core_ext/numeric/bytes"
 require "relation_to_struct"
 require "ruby2_keywords"
 
@@ -31,6 +32,17 @@ def self.configure
 
   LOCK_TIMEOUT_SECONDS = 5
   LOCK_FAILURE_RETRY_DELAY_MULTLIPLIER = 5
+  SMALL_TABLE_THRESHOLD_BYTES = 10.megabytes
+
+  PARTITION_TYPES = %i[range list hash]
+
+  PARTMAN_UPDATE_CONFIG_OPTIONS = %i[
+    infinite_time_partitions
+    inherit_privileges
+    premake
+    retention
+    retention_keep_table
+  ]
 
   # Safe versus unsafe in this context specifically means the following:
   # - Safe operations will not block for long periods of time.

lib/pg_ha_migrations/relation.rb

@@ -96,6 +96,19 @@ def partitions(include_sub_partitions: false, include_self: false)
 
       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

lib/pg_ha_migrations/safe_statements.rb

@@ -1,14 +1,4 @@
 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
@@ -154,6 +144,22 @@ def unsafe_make_column_not_nullable(table, column, options={}) # options arg is
     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
@@ -307,8 +313,8 @@ def unsafe_remove_constraint(table, name:)
   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
@@ -447,7 +453,7 @@ def safe_partman_update_config(table, **options)
   end
 
   def unsafe_partman_update_config(table, **options)
-    invalid_options = options.keys - PARTMAN_UPDATE_CONFIG_OPTIONS
+    invalid_options = options.keys - PgHaMigrations::PARTMAN_UPDATE_CONFIG_OPTIONS
 
     raise ArgumentError, "Unrecognized argument(s): #{invalid_options}" unless invalid_options.empty?
 
@@ -631,4 +637,16 @@ def adjust_statement_timeout(timeout_seconds, &block)
       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