DEV: Add protection for migrations that creates index concurrently (#31763)

This commit updates `Migration::SafeMigrate` to protect against unsafe
ways of adding a Postgres index concurrently.

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 .

Therefore, the simplest way for us to ensure that migrations that create
indexes concurrently are idempotent is to follow postgres'
recommendation of dropping the index first before trying to create the
index concurrently.

lib/migration/safe_migrate.rb

@@ -77,6 +77,8 @@ class Migration::SafeMigrate
     return if PG::Connection.method_defined?(:exec_migrator_unpatched)
     return if ENV["RAILS_ENV"] == "production"
 
+    @@migration_sqls = []
+
     PG::Connection.class_eval do
       alias_method :exec_migrator_unpatched, :exec
       alias_method :async_exec_migrator_unpatched, :async_exec
@@ -97,6 +99,8 @@ class Migration::SafeMigrate
     return if !PG::Connection.method_defined?(:exec_migrator_unpatched)
     return if ENV["RAILS_ENV"] == "production"
 
+    @@migration_sqls = []
+
     PG::Connection.class_eval do
       alias_method :exec, :exec_migrator_unpatched
       alias_method :async_exec, :async_exec_migrator_unpatched
@@ -116,7 +120,24 @@ 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
+
     if sql =~ /\A\s*(?:drop\s+table|alter\s+table.*rename\s+to)\s+/i
       $stdout.puts("", <<~TEXT)
         WARNING
@@ -147,6 +168,16 @@ class Migration::SafeMigrate
         in use by live applications.
       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
+      match = sql.match(/\bON\s+(?:ONLY\s+)?(?:"([^"]+)"|([a-zA-Z0-9_\.]+))/i)
+      table_name = match[1] || match[2]
+
+      if !@@migration_sqls.any? { |migration_sql|
+           migration_sql =~ /\A\s*drop\s+index\s+(?:concurrently\s+)?if\s+exists\s+/i &&
+             migration_sql.include?(table_name)
+         }
+        $stdout.puts("", UNSAFE_DROP_INDEX_CONCURRENTLY_WARNING)
+      end
     end
   end
 

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

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+class CreateIndexConcurrently < ActiveRecord::Migration[7.2]
+  disable_ddl_transaction!
+
+  def up
+    execute <<~SQL
+    DROP INDEX IF EXISTS some_notifications_index;
+    SQL
+
+    execute <<~SQL
+    CREATE INDEX CONCURRENTLY some_notifications_index ON notifications(user_id);
+    SQL
+
+    execute <<~SQL
+    DROP INDEX some_notifications_index;
+    SQL
+  end
+
+  def down
+    raise "not tested"
+  end
+end

spec/fixtures/db/migrate/create_index_concurrently_unsafe/20230309014015_create_index_concurrently.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+class CreateIndexConcurrently < ActiveRecord::Migration[7.2]
+  disable_ddl_transaction!
+
+  def up
+    execute <<~SQL
+    DROP INDEX IF EXISTS some_wrong_index;
+    SQL
+
+    execute <<~SQL
+    CREATE INDEX CONCURRENTLY some_notifications_index ON notifications(user_id, id DESC, read, topic_id);
+    SQL
+  end
+
+  def down
+    raise "not tested"
+  end
+end

spec/lib/migration/safe_migrate_spec.rb

@@ -71,6 +71,24 @@ 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
+    path = File.expand_path "#{Rails.root}/spec/fixtures/db/migrate/create_index_concurrently_safe"
+    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)
+  end
+
+  it "bans running a migration that creates an index concurrently without first dropping the index if it exists" do
+    Migration::SafeMigrate.enable!
+
+    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) }
+
+    expect(output).to include(described_class::UNSAFE_DROP_INDEX_CONCURRENTLY_WARNING)
+  end
+
   it "allows dropping NOT NULL" do
     Migration::SafeMigrate.enable!