DEV: Protection for migrations that creates index concurrently take 2 (#31792)

This is a follow up to 6820622467ab3613e824f0cb6219def2a575bc1d.

This commit addresses migrations that uses `remove_index` with the
`if_exists: true` option to drop an existing index before creating an
index using the `concurrently` option.

This commit also reruns two migration which may have caused indexes to
be left in an `invalid` state.

### Reviewers Note

Plugin tests are failing due to
https://github.com/discourse/discourse-translator/pull/251

db/migrate/20241115143259_add_index_to_users_ip_address.rb

@@ -1,9 +0,0 @@
-# frozen_string_literal: true
-
-class AddIndexToUsersIpAddress < ActiveRecord::Migration[7.1]
-  disable_ddl_transaction!
-
-  def change
-    add_index :users, :ip_address, algorithm: :concurrently, name: "idx_users_ip_address"
-  end
-end

db/migrate/20250313044812_remove_unique_constraint_from_tag_users_indexes.rb

@@ -1,10 +1,21 @@
 # frozen_string_literal: true
+
 class RemoveUniqueConstraintFromTagUsersIndexes < ActiveRecord::Migration[7.1]
   disable_ddl_transaction!
 
   def up
-    remove_index :tag_users, name: :idx_tag_users_ix1, algorithm: :concurrently
+    remove_index :tag_users, name: :idx_tag_users_ix1, algorithm: :concurrently, if_exists: true
-    remove_index :tag_users, name: :idx_tag_users_ix2, algorithm: :concurrently
+    remove_index :tag_users, name: :idx_tag_users_ix2, algorithm: :concurrently, if_exists: true
+
+    remove_index :tag_users,
+                 %i[user_id tag_id notification_level],
+                 algorithm: :concurrently,
+                 if_exists: true
+
+    remove_index :tag_users,
+                 %i[tag_id user_id notification_level],
+                 algorithm: :concurrently,
+                 if_exists: true
 
     add_index :tag_users, %i[user_id tag_id notification_level], algorithm: :concurrently
     add_index :tag_users, %i[tag_id user_id notification_level], algorithm: :concurrently

db/migrate/20250313045010_add_index_to_users_ip_address.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+class AddIndexToUsersIpAddress < ActiveRecord::Migration[7.1]
+  disable_ddl_transaction!
+
+  def up
+    remove_index :users,
+                 :ip_address,
+                 algorithm: :concurrently,
+                 name: "idx_users_ip_address",
+                 if_exists: true
+
+    add_index :users, :ip_address, algorithm: :concurrently, name: "idx_users_ip_address"
+  end
+
+  def down
+    raise ActiveRecord::IrreversibleMigration
+  end
+end

lib/migration/safe_migrate.rb

@@ -78,6 +78,20 @@ class Migration::SafeMigrate
     return if ENV["RAILS_ENV"] == "production"
 
     @@migration_sqls = []
+    @@activerecord_remove_indexes = []
+
+    ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.class_eval do
+      alias_method :original_remove_index, :remove_index
+
+      def remove_index(table_name, column_name = nil, **options)
+        @@activerecord_remove_indexes << (
+          options[:name] || index_name(table_name, options.merge(column: column_name))
+        )
+
+        # Call the original method
+        original_remove_index(table_name, column_name, **options)
+      end
+    end
 
     PG::Connection.class_eval do
       alias_method :exec_migrator_unpatched, :exec
@@ -99,7 +113,13 @@ class Migration::SafeMigrate
     return if !PG::Connection.method_defined?(:exec_migrator_unpatched)
     return if ENV["RAILS_ENV"] == "production"
 
-    @@migration_sqls = []
+    @@migration_sqls.clear
+    @@activerecord_remove_indexes.clear
+
+    ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.class_eval do
+      alias_method :remove_index, :original_remove_index
+      remove_method :original_remove_index
+    end
 
     PG::Connection.class_eval do
       alias_method :exec, :exec_migrator_unpatched
@@ -120,21 +140,6 @@ class Migration::SafeMigrate
     end
   end
 
-  UNSAFE_DROP_INDEX_CONCURRENTLY_WARNING = <<~TEXT
-  WARNING
-  -------------------------------------------------------------------------------------
-  An attempt was made to create an index concurrently in a migration without first dropping the index.
-
-  Per postgres documentation:
-
-    If a problem arises while scanning the table, such as a deadlock or a uniqueness violation in a unique index,
-    the CREATE INDEX command will fail but leave behind an “invalid” index. This index will be ignored for querying
-    purposes because it might be incomplete; however it will still consume update overhead. The recommended recovery
-    method in such cases is to drop the index and try again to perform CREATE INDEX CONCURRENTLY .
-
-  Please update the migration to first drop the index if it exists before creating it concurrently.
-  TEXT
-
   def self.protect!(sql)
     @@migration_sqls << sql
 
@@ -169,15 +174,37 @@ class Migration::SafeMigrate
       TEXT
       raise Discourse::InvalidMigration, "Attempt was made to rename or delete column"
     elsif sql =~ /\A\s*create\s+(?:unique\s+)?index\s+concurrently\s+/i
+      index_name =
+        sql.match(/\bINDEX\s+CONCURRENTLY\s+(?:IF\s+NOT\s+EXISTS\s+)?"?([a-zA-Z0-9_\.]+)"?/i)[1]
+
+      return if @@activerecord_remove_indexes.include?(index_name)
+
       match = sql.match(/\bON\s+(?:ONLY\s+)?(?:"([^"]+)"|([a-zA-Z0-9_\.]+))/i)
       table_name = match[1] || match[2]
 
-      if !@@migration_sqls.any? { |migration_sql|
+      has_drop_index_statement =
-           migration_sql =~ /\A\s*drop\s+index\s+(?:concurrently\s+)?if\s+exists\s+/i &&
+        @@migration_sqls.any? do |migration_sql|
-             migration_sql.include?(table_name)
+          migration_sql =~ /\A\s*drop\s+index/i && migration_sql.include?(table_name) &&
-         }
+            migration_sql.include?(index_name)
-        $stdout.puts("", UNSAFE_DROP_INDEX_CONCURRENTLY_WARNING)
+        end
-      end
+
+      return if has_drop_index_statement
+
+      raise(Discourse::InvalidMigration, <<~RAW)
+      WARNING
+      -------------------------------------------------------------------------------------
+      An attempt was made to create an index concurrently in a migration without first dropping the index.
+      SQL used was: '#{sql}'
+
+      Per postgres documentation:
+
+        If a problem arises while scanning the table, such as a deadlock or a uniqueness violation in a unique index,
+        the CREATE INDEX command will fail but leave behind an “invalid” index. This index will be ignored for querying
+        purposes because it might be incomplete; however it will still consume update overhead. The recommended recovery
+        method in such cases is to drop the index and try again to perform CREATE INDEX CONCURRENTLY .
+
+      Please update the migration to first drop the index if it exists before creating it concurrently.
+      RAW
     end
   end
 

spec/fixtures/db/migrate/create_index_concurrently_safe/20230309014016_create_index_concurrently.rb

@@ -11,10 +11,6 @@ class CreateIndexConcurrently < ActiveRecord::Migration[7.2]
     execute <<~SQL
     CREATE INDEX CONCURRENTLY some_notifications_index ON notifications(user_id);
     SQL
-
-    execute <<~SQL
-    DROP INDEX some_notifications_index;
-    SQL
   end
 
   def down

spec/fixtures/db/migrate/create_index_concurrently_safe_activerecord/20230309014016_some_migration.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+class SomeMigration < ActiveRecord::Migration[7.2]
+  disable_ddl_transaction!
+
+  def up
+    remove_index :notifications, %i[user_id id], if_exists: true
+    add_index :notifications, %i[user_id id], algorithm: :concurrently
+  end
+
+  def down
+    raise "not tested"
+  end
+end

spec/lib/migration/safe_migrate_spec.rb

@@ -71,11 +71,44 @@ RSpec.describe Migration::SafeMigrate do
     expect { User.first.username }.not_to raise_error
   end
 
-  it "allows running a migration that creates an index concurrently if it drops the index first" do
+  it "allows running a migration that creates an index concurrently if it checks if the index exists first" do
-    path = File.expand_path "#{Rails.root}/spec/fixtures/db/migrate/create_index_concurrently_safe"
+    Migration::SafeMigrate.enable!
-    output = capture_stdout { expect do migrate_up(path) end.to raise_error(StandardError) }
 
-    expect(output).not_to include(described_class::UNSAFE_DROP_INDEX_CONCURRENTLY_WARNING)
+    path =
+      File.expand_path "#{Rails.root}/spec/fixtures/db/migrate/create_index_concurrently_safe_activerecord"
+
+    error = nil
+
+    capture_stdout do
+      begin
+        migrate_up(path)
+      rescue StandardError => e
+        error = e
+      end
+    end
+
+    expect(error.cause.cause.message).to include(
+      "CREATE INDEX CONCURRENTLY cannot run inside a transaction block",
+    )
+  end
+
+  it "allows running a migration that creates an index concurrently if it drops the index first" do
+    Migration::SafeMigrate.enable!
+
+    path = File.expand_path "#{Rails.root}/spec/fixtures/db/migrate/create_index_concurrently_safe"
+    error = nil
+
+    capture_stdout do
+      begin
+        migrate_up(path)
+      rescue StandardError => e
+        error = e
+      end
+    end
+
+    expect(error.cause.cause.message).to include(
+      "CREATE INDEX CONCURRENTLY cannot run inside a transaction block",
+    )
   end
 
   it "bans running a migration that creates an index concurrently without first dropping the index if it exists" do
@@ -84,9 +117,17 @@ RSpec.describe Migration::SafeMigrate do
     path =
       File.expand_path("#{Rails.root}/spec/fixtures/db/migrate/create_index_concurrently_unsafe")
 
-    output = capture_stdout { expect do migrate_up(path) end.to raise_error(StandardError) }
+    error = nil
 
-    expect(output).to include(described_class::UNSAFE_DROP_INDEX_CONCURRENTLY_WARNING)
+    capture_stdout do
+      migrate_up(path)
+    rescue StandardError => e
+      error = e
+    end
+
+    expect(error.message).to include(
+      "An attempt was made to create an index concurrently in a migration without first dropping the index.",
+    )
   end
 
   it "allows dropping NOT NULL" do