add check for table size on disk

README.md

@@ -484,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

@@ -31,6 +31,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
@@ -160,11 +150,11 @@ def safe_add_index_on_empty_table(table, columns, options={})
     end
 
     # Avoids taking out an unnecessary SHARE lock if the table does have data
-    _ensure_empty_table!(table)
+    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_empty_table!(table)
+      ensure_small_table!(table, empty: true)
 
       unsafe_add_index(table, columns, **options)
     end
@@ -323,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
@@ -463,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?
 
@@ -525,14 +515,6 @@ def _type_is_enum(type)
     ActiveRecord::Base.connection.select_values("SELECT typname FROM pg_type JOIN pg_enum ON pg_type.oid = pg_enum.enumtypid").include?(type.to_s)
   end
 
-  def _ensure_empty_table!(table)
-    table = PgHaMigrations::Table.from_table_name(table)
-
-    if connection.select_value("SELECT EXISTS (SELECT 1 FROM #{table.fully_qualified_name} LIMIT 1)")
-      raise PgHaMigrations::InvalidMigrationError, "Table #{table.inspect} has rows. Please use safe_add_concurrent_index instead."
-    end
-  end
-
   def migrate(direction)
     if respond_to?(:change)
       raise PgHaMigrations::UnsupportedMigrationError, "Tracking changes for automated rollback is not supported; use explicit #up instead."
@@ -655,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

spec/safe_statements_spec.rb

@@ -1319,10 +1319,40 @@ def up
 
             allow(ActiveRecord::Base.connection).to receive(:select_value).and_call_original
             expect(ActiveRecord::Base.connection).to receive(:select_value).with(/SELECT EXISTS/).once.and_call_original
+            expect(ActiveRecord::Base.connection).to_not receive(:select_value).with(/pg_total_relation_size/)
 
             expect do
               test_migration.suppress_messages { test_migration.migrate(:up) }
-            end.to raise_error(PgHaMigrations::InvalidMigrationError, "Table \"foos\" has rows. Please use safe_add_concurrent_index instead.")
+            end.to raise_error(PgHaMigrations::InvalidMigrationError, "Table \"foos\" has rows")
+
+            indexes = ActiveRecord::Base.connection.indexes("foos")
+            expect(indexes).to be_empty
+          end
+
+          it "raises error when table is larger than small table threshold" do
+            setup_migration = Class.new(migration_klass) do
+              def up
+                unsafe_create_table :foos
+              end
+            end
+
+            setup_migration.suppress_messages { setup_migration.migrate(:up) }
+
+            test_migration = Class.new(migration_klass) do
+              def up
+                safe_add_index_on_empty_table :foos, [:bar]
+              end
+            end
+
+            stub_const("PgHaMigrations::SMALL_TABLE_THRESHOLD_BYTES", 1.kilobyte)
+
+            allow(ActiveRecord::Base.connection).to receive(:select_value).and_call_original
+            expect(ActiveRecord::Base.connection).to receive(:select_value).with(/SELECT EXISTS/).once.and_call_original
+            expect(ActiveRecord::Base.connection).to receive(:select_value).with(/pg_total_relation_size/).once.and_call_original
+
+            expect do
+              test_migration.suppress_messages { test_migration.migrate(:up) }
+            end.to raise_error(PgHaMigrations::InvalidMigrationError, "Table \"foos\" is larger than 1024 bytes")
 
             indexes = ActiveRecord::Base.connection.indexes("foos")
             expect(indexes).to be_empty
@@ -1345,6 +1375,7 @@ def up
             end
 
             allow(ActiveRecord::Base.connection).to receive(:select_value).and_call_original
+            expect(ActiveRecord::Base.connection).to receive(:select_value).with(/pg_total_relation_size/).once.and_call_original
             expect(ActiveRecord::Base.connection).to receive(:select_value).with(/SELECT EXISTS/).twice.and_wrap_original do |m, *args|
               m.call(*args).tap do
                 ActiveRecord::Base.connection.execute("INSERT INTO foos DEFAULT VALUES")
@@ -1353,7 +1384,7 @@ def up
 
             expect do
               test_migration.suppress_messages { test_migration.migrate(:up) }
-            end.to raise_error(PgHaMigrations::InvalidMigrationError, "Table \"foos\" has rows. Please use safe_add_concurrent_index instead.")
+            end.to raise_error(PgHaMigrations::InvalidMigrationError, "Table \"foos\" has rows")
 
             indexes = ActiveRecord::Base.connection.indexes("foos")
             expect(indexes).to be_empty
@@ -4415,6 +4446,134 @@ def up
           end
         end
 
+        describe "ensure_small_table!" do
+          it "does not raise error when empty: false and table is below threshold and has rows" do
+            setup_migration = Class.new(migration_klass) do
+              def up
+                safe_create_table :foos
+
+                unsafe_execute "INSERT INTO foos DEFAULT VALUES"
+              end
+            end
+
+            setup_migration.suppress_messages { setup_migration.migrate(:up) }
+
+            test_migration = Class.new(migration_klass) do
+              def up
+                ensure_small_table! :foos
+              end
+            end
+
+            allow(ActiveRecord::Base.connection).to receive(:select_value).and_call_original
+            expect(ActiveRecord::Base.connection).to_not receive(:select_value).with(/SELECT EXISTS/)
+            expect(ActiveRecord::Base.connection).to receive(:select_value).with(/pg_total_relation_size/).once.and_call_original
+
+            expect do
+              test_migration.suppress_messages { test_migration.migrate(:up) }
+            end.to_not raise_error
+          end
+
+          it "does not raise error when empty: true and table is below threshold and is empty" do
+            setup_migration = Class.new(migration_klass) do
+              def up
+                safe_create_table :foos
+              end
+            end
+
+            setup_migration.suppress_messages { setup_migration.migrate(:up) }
+
+            test_migration = Class.new(migration_klass) do
+              def up
+                ensure_small_table! :foos, empty: true
+              end
+            end
+
+            allow(ActiveRecord::Base.connection).to receive(:select_value).and_call_original
+            expect(ActiveRecord::Base.connection).to receive(:select_value).with(/SELECT EXISTS/).once.and_call_original
+            expect(ActiveRecord::Base.connection).to receive(:select_value).with(/pg_total_relation_size/).once.and_call_original
+
+            expect do
+              test_migration.suppress_messages { test_migration.migrate(:up) }
+            end.to_not raise_error
+          end
+
+          it "raises error when empty: true and table has rows" do
+            setup_migration = Class.new(migration_klass) do
+              def up
+                safe_create_table :foos
+
+                unsafe_execute "INSERT INTO foos DEFAULT VALUES"
+              end
+            end
+
+            setup_migration.suppress_messages { setup_migration.migrate(:up) }
+
+            test_migration = Class.new(migration_klass) do
+              def up
+                ensure_small_table! :foos, empty: true
+              end
+            end
+
+            allow(ActiveRecord::Base.connection).to receive(:select_value).and_call_original
+            expect(ActiveRecord::Base.connection).to receive(:select_value).with(/SELECT EXISTS/).once.and_call_original
+            expect(ActiveRecord::Base.connection).to_not receive(:select_value).with(/pg_total_relation_size/)
+
+            expect do
+              test_migration.suppress_messages { test_migration.migrate(:up) }
+            end.to raise_error(PgHaMigrations::InvalidMigrationError, "Table \"foos\" has rows")
+          end
+
+          it "raises error when empty: true and table is above threshold and is empty" do
+            setup_migration = Class.new(migration_klass) do
+              def up
+                safe_create_table :foos
+              end
+            end
+
+            setup_migration.suppress_messages { setup_migration.migrate(:up) }
+
+            test_migration = Class.new(migration_klass) do
+              def up
+                ensure_small_table! :foos, empty: true, threshold: 1.kilobyte
+              end
+            end
+
+            allow(ActiveRecord::Base.connection).to receive(:select_value).and_call_original
+            expect(ActiveRecord::Base.connection).to receive(:select_value).with(/SELECT EXISTS/).once.and_call_original
+            expect(ActiveRecord::Base.connection).to receive(:select_value).with(/pg_total_relation_size/).once.and_call_original
+
+            expect do
+              test_migration.suppress_messages { test_migration.migrate(:up) }
+            end.to raise_error(PgHaMigrations::InvalidMigrationError, "Table \"foos\" is larger than 1024 bytes")
+          end
+
+          it "raises error when empty: false and table is above threshold and has rows" do
+            setup_migration = Class.new(migration_klass) do
+              def up
+                safe_create_table :foos
+
+                unsafe_execute "INSERT INTO foos DEFAULT VALUES"
+              end
+            end
+
+            setup_migration.suppress_messages { setup_migration.migrate(:up) }
+
+            test_migration = Class.new(migration_klass) do
+              def up
+                ensure_small_table! :foos, threshold: 1.kilobyte
+              end
+            end
+
+            allow(ActiveRecord::Base.connection).to receive(:select_value).and_call_original
+            expect(ActiveRecord::Base.connection).to_not receive(:select_value).with(/SELECT EXISTS/)
+            expect(ActiveRecord::Base.connection).to receive(:select_value).with(/pg_total_relation_size/).once.and_call_original
+
+            expect do
+              test_migration.suppress_messages { test_migration.migrate(:up) }
+            end.to raise_error(PgHaMigrations::InvalidMigrationError, "Table \"foos\" is larger than 1024 bytes")
+          end
+        end
+
         describe "unsafe transformations" do
           it "renames create_table to unsafe_create_table" do
             migration = Class.new(migration_klass) do