code review comments

Co-authored-by: Ryan Krage <[email protected]>

lib/pg_ha_migrations/safe_statements.rb

@@ -155,13 +155,15 @@ def unsafe_make_column_not_nullable(table, column, options={}) # options arg is
   end
 
   def safe_add_index_on_empty_table(table, columns, options={})
-    table_struct = PgHaMigrations::Table.from_table_name(table)
+    # Avoids taking out an unnecessary SHARE lock if the table does have data
+    _ensure_empty_table!(table)
 
-    if connection.select_value("SELECT EXISTS (SELECT 1 FROM #{table_struct.fully_qualified_name} LIMIT 1)")
-      raise PgHaMigrations::InvalidMigrationError, "Table #{table_struct.inspect} has rows. Please use safe_add_concurrent_index instead."
-    end
+    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)
 
-    unsafe_add_index(table, columns, **options)
+      unsafe_add_index(table, columns, **options)
+    end
   end
 
   def safe_add_concurrent_index(table, columns, options={})
@@ -214,18 +216,21 @@ def safe_add_concurrent_partitioned_index(
       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,
-      where: where,
-      comment: comment,
-      algorithm: :only, # see lib/pg_ha_migrations/hacks/add_index_on_only.rb
-    )
+    # 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|
       # CREATE INDEX CONCURRENTLY ON child_table
@@ -513,6 +518,14 @@ 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."

lib/pg_ha_migrations/unsafe_statements.rb

@@ -80,15 +80,7 @@ def unsafe_add_index(table, column_names, options = {})
       raise PgHaMigrations::InvalidMigrationError, "ActiveRecord drops the :opclass option when supplying a string containing an expression or list of columns; instead either supply an array of columns or include the opclass in the string for each column"
     end
 
-    if options[:algorithm] == :concurrently
-      # Cannot be run inside a transaction so we are unable to use
-      # safely_acquire_lock_for_table to take out ShareUpdateExclusiveLock
-      execute_ancestor_statement(:add_index, table, column_names, **options)
-    else
-      safely_acquire_lock_for_table(table, mode: :share) do
-        execute_ancestor_statement(:add_index, table, column_names, **options)
-      end
-    end
+    execute_ancestor_statement(:add_index, table, column_names, **options)
   end
 
   ruby2_keywords def execute_ancestor_statement(method_name, *args, &block)

spec/safe_statements_spec.rb

@@ -1300,6 +1300,40 @@ def up
               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 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.")
+
+            indexes = ActiveRecord::Base.connection.indexes("foos")
+            expect(indexes).to be_empty
+          end
+
+          it "raises error when table receives writes immediately after the first check" do
+            setup_migration = Class.new(migration_klass) do
+              def up
+                unsafe_create_table :foos
+                unsafe_add_column :foos, :bar, :text
+              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
+
+            allow(ActiveRecord::Base.connection).to receive(:select_value).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")
+              end
+            end
+
             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.")