Remove default connection for WithLockRetries

Make the :connection argument for WithLockRetries, by removing the
default of ActiveRecord::Base. This prevents accidental use of the wrong
connection once we begin using multiple databases.

.rubocop_todo/database/multiple_databases.yml

@@ -19,7 +19,6 @@ Database/MultipleDatabases:
   - lib/gitlab/database/migrations/observers/query_log.rb
   - lib/gitlab/database/partitioning_migration_helpers/backfill_partitioned_table.rb
   - lib/gitlab/database.rb
-  - lib/gitlab/database/with_lock_retries.rb
   - lib/gitlab/gitlab_import/importer.rb
   - lib/gitlab/health_checks/db_check.rb
   - lib/gitlab/import_export/base/relation_factory.rb

db/migrate/20210316171009_create_packages_helm_file_metadata.rb

@@ -5,7 +5,7 @@ class CreatePackagesHelmFileMetadata < ActiveRecord::Migration[6.0]
 
   DOWNTIME = false
 
-  def change
+  def up
     create_table_with_constraints :packages_helm_file_metadata, id: false do |t|
       t.timestamps_with_timezone
       t.references :package_file, primary_key: true, index: false, default: nil, null: false, foreign_key: { to_table: :packages_package_files, on_delete: :cascade }, type: :bigint
@@ -17,4 +17,10 @@ class CreatePackagesHelmFileMetadata < ActiveRecord::Migration[6.0]
       t.index :channel
     end
   end
+
+  def down
+    with_lock_retries do
+      drop_table :packages_helm_file_metadata
+    end
+  end
 end

db/migrate/20210504153354_create_clusters_integration_elasticstack.rb

@@ -3,7 +3,7 @@
 class CreateClustersIntegrationElasticstack < ActiveRecord::Migration[6.0]
   include Gitlab::Database::MigrationHelpers
 
-  def change
+  def up
     create_table_with_constraints :clusters_integration_elasticstack, id: false do |t|
       t.timestamps_with_timezone null: false
       t.references :cluster, primary_key: true, default: nil, index: false, foreign_key: { on_delete: :cascade }
@@ -12,4 +12,10 @@ class CreateClustersIntegrationElasticstack < ActiveRecord::Migration[6.0]
       t.text_limit :chart_version, 10
     end
   end
+
+  def down
+    with_lock_retries do
+      drop_table :clusters_integration_elasticstack
+    end
+  end
 end

db/migrate/20210628154900_create_detached_partitions_table.rb

@@ -3,7 +3,7 @@
 class CreateDetachedPartitionsTable < ActiveRecord::Migration[6.1]
   include Gitlab::Database::MigrationHelpers
 
-  def change
+  def up
     create_table_with_constraints :detached_partitions do |t|
       t.timestamps_with_timezone null: false
       t.datetime_with_timezone :drop_after, null: false
@@ -14,4 +14,10 @@ class CreateDetachedPartitionsTable < ActiveRecord::Migration[6.1]
       t.text_limit :table_name, 63
     end
   end
+
+  def down
+    with_lock_retries do
+      drop_table :detached_partitions
+    end
+  end
 end

db/migrate/20210720140841_create_postgres_async_indexes_table.rb

@@ -3,7 +3,7 @@
 class CreatePostgresAsyncIndexesTable < ActiveRecord::Migration[6.1]
   include Gitlab::Database::MigrationHelpers
 
-  def change
+  def up
     create_table_with_constraints :postgres_async_indexes do |t|
       t.timestamps_with_timezone null: false
 
@@ -18,4 +18,10 @@ class CreatePostgresAsyncIndexesTable < ActiveRecord::Migration[6.1]
       t.index :name, unique: true
     end
   end
+
+  def down
+    with_lock_retries do
+      drop_table :postgres_async_indexes
+    end
+  end
 end

db/migrate/20210729081351_create_topics.rb

@@ -3,7 +3,7 @@
 class CreateTopics < ActiveRecord::Migration[6.1]
   include Gitlab::Database::MigrationHelpers
 
-  def change
+  def up
     create_table_with_constraints :topics do |t|
       t.text :name, null: false
       t.text_limit :name, 255
@@ -13,4 +13,10 @@ class CreateTopics < ActiveRecord::Migration[6.1]
       t.timestamps_with_timezone
     end
   end
+
+  def down
+    with_lock_retries do
+      drop_table :topics
+    end
+  end
 end

lib/gitlab/database/migration_helpers.rb

@@ -429,6 +429,7 @@ module Gitlab
       def with_lock_retries(*args, **kwargs, &block)
         raise_on_exhaustion = !!kwargs.delete(:raise_on_exhaustion)
         merged_args = {
+          connection: connection,
           klass: self.class,
           logger: Gitlab::BackgroundMigration::Logger,
           allow_savepoints: true

lib/gitlab/database/migrations/lock_retry_mixin.rb

@@ -9,6 +9,10 @@ module Gitlab
             migration.class
           end
 
+          def migration_connection
+            migration.connection
+          end
+
           def enable_lock_retries?
             # regular AR migrations don't have this,
             # only ones inheriting from Gitlab::Database::Migration have
@@ -24,6 +28,7 @@ module Gitlab
           def ddl_transaction(migration, &block)
             if use_transaction?(migration) && migration.enable_lock_retries?
               Gitlab::Database::WithLockRetries.new(
+                connection: migration.migration_connection,
                 klass: migration.migration_class,
                 logger: Gitlab::BackgroundMigration::Logger
               ).run(raise_on_exhaustion: false, &block)

lib/gitlab/database/schema_helpers.rb

@@ -73,6 +73,7 @@ module Gitlab
 
       def with_lock_retries(&block)
         Gitlab::Database::WithLockRetries.new(
+          connection: connection,
           klass: self.class,
           logger: Gitlab::BackgroundMigration::Logger
         ).run(&block)

lib/gitlab/database/with_lock_retries.rb

@@ -61,7 +61,7 @@ module Gitlab
         [10.seconds, 10.minutes]
       ].freeze
 
-      def initialize(logger: NULL_LOGGER, allow_savepoints: true, timing_configuration: DEFAULT_TIMING_CONFIGURATION, klass: nil, env: ENV, connection: ActiveRecord::Base.connection)
+      def initialize(logger: NULL_LOGGER, allow_savepoints: true, timing_configuration: DEFAULT_TIMING_CONFIGURATION, klass: nil, env: ENV, connection:)
         @logger = logger
         @klass = klass
         @allow_savepoints = allow_savepoints

spec/lib/gitlab/database/migrations/lock_retry_mixin_spec.rb

@@ -3,7 +3,8 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Migrations::LockRetryMixin do
   describe Gitlab::Database::Migrations::LockRetryMixin::ActiveRecordMigrationProxyLockRetries do
-    let(:migration) { double }
+    let(:connection) { ActiveRecord::Base.connection }
+    let(:migration) { double(connection: connection) }
     let(:return_value) { double }
     let(:class_def) do
       Class.new do
@@ -40,6 +41,18 @@ RSpec.describe Gitlab::Database::Migrations::LockRetryMixin do
         expect(result).to eq(return_value)
       end
     end
+
+    describe '#migration_connection' do
+      subject { class_def.new(migration).migration_connection }
+
+      it 'retrieves actual migration connection from #migration' do
+        expect(migration).to receive(:connection).and_return(return_value)
+
+        result = subject
+
+        expect(result).to eq(return_value)
+      end
+    end
   end
 
   describe Gitlab::Database::Migrations::LockRetryMixin::ActiveRecordMigratorLockRetries do
@@ -96,7 +109,8 @@ RSpec.describe Gitlab::Database::Migrations::LockRetryMixin do
 
     context 'with transactions enabled and lock retries enabled' do
       let(:receiver) { double('receiver', use_transaction?: true)}
-      let(:migration) { double('migration', enable_lock_retries?: true) }
+      let(:migration) { double('migration', migration_connection: connection, enable_lock_retries?: true) }
+      let(:connection) { ActiveRecord::Base.connection }
 
       it 'calls super method' do
         p = proc { }

spec/lib/gitlab/database/with_lock_retries_outside_transaction_spec.rb

@@ -5,7 +5,8 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
   let(:env) { {} }
   let(:logger) { Gitlab::Database::WithLockRetries::NULL_LOGGER }
-  let(:subject) { described_class.new(env: env, logger: logger, timing_configuration: timing_configuration) }
+  let(:subject) { described_class.new(connection: connection, env: env, logger: logger, timing_configuration: timing_configuration) }
+  let(:connection) { ActiveRecord::Base.retrieve_connection }
 
   let(:timing_configuration) do
     [
@@ -67,7 +68,7 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
             WHERE t.relkind = 'r' AND l.mode = 'ExclusiveLock' AND t.relname = '#{Project.table_name}'
           """
 
-          expect(ActiveRecord::Base.connection.execute(check_exclusive_lock_query).to_a).to be_present
+          expect(connection.execute(check_exclusive_lock_query).to_a).to be_present
         end
       end
 
@@ -96,8 +97,8 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
               lock_fiber.resume
             end
 
-            ActiveRecord::Base.transaction do
-              ActiveRecord::Base.connection.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
+            connection.transaction do
+              connection.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
               lock_acquired = true
             end
           end
@@ -115,7 +116,7 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
         context 'setting the idle transaction timeout' do
           context 'when there is no outer transaction: disable_ddl_transaction! is set in the migration' do
             it 'does not disable the idle transaction timeout' do
-              allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
+              allow(connection).to receive(:transaction_open?).and_return(false)
               allow(subject).to receive(:run_block_with_lock_timeout).once.and_raise(ActiveRecord::LockWaitTimeout)
               allow(subject).to receive(:run_block_with_lock_timeout).once
 
@@ -127,7 +128,7 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
 
           context 'when there is outer transaction: disable_ddl_transaction! is not set in the migration' do
             it 'disables the idle transaction timeout so the code can sleep and retry' do
-              allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(true)
+              allow(connection).to receive(:transaction_open?).and_return(true)
 
               n = 0
               allow(subject).to receive(:run_block_with_lock_timeout).twice do
@@ -184,8 +185,8 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
             subject.run(raise_on_exhaustion: true) do
               lock_attempts += 1
 
-              ActiveRecord::Base.transaction do
-                ActiveRecord::Base.connection.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
+              connection.transaction do
+                connection.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
                 lock_acquired = true
               end
             end
@@ -199,11 +200,11 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
       context 'when statement timeout is reached' do
         it 'raises StatementInvalid error' do
           lock_acquired = false
-          ActiveRecord::Base.connection.execute("SET statement_timeout='100ms'")
+          connection.execute("SET statement_timeout='100ms'")
 
           expect do
             subject.run do
-              ActiveRecord::Base.connection.execute("SELECT 1 FROM pg_sleep(0.11)") # 110ms
+              connection.execute("SELECT 1 FROM pg_sleep(0.11)") # 110ms
               lock_acquired = true
             end
           end.to raise_error(ActiveRecord::StatementInvalid)
@@ -216,11 +217,11 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
 
   context 'restore local database variables' do
     it do
-      expect { subject.run {} }.not_to change { ActiveRecord::Base.connection.execute("SHOW lock_timeout").to_a }
+      expect { subject.run {} }.not_to change { connection.execute("SHOW lock_timeout").to_a }
     end
 
     it do
-      expect { subject.run {} }.not_to change { ActiveRecord::Base.connection.execute("SHOW idle_in_transaction_session_timeout").to_a }
+      expect { subject.run {} }.not_to change { connection.execute("SHOW idle_in_transaction_session_timeout").to_a }
     end
   end
 
@@ -228,8 +229,8 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
     let(:timing_configuration) { [[0.015.seconds, 0.025.seconds], [0.015.seconds, 0.025.seconds]] } # 15ms, 25ms
 
     it 'executes `SET lock_timeout` using the configured timeout value in milliseconds' do
-      expect(ActiveRecord::Base.connection).to receive(:execute).with('RESET idle_in_transaction_session_timeout; RESET lock_timeout').and_call_original
-      expect(ActiveRecord::Base.connection).to receive(:execute).with("SET lock_timeout TO '15ms'").and_call_original
+      expect(connection).to receive(:execute).with('RESET idle_in_transaction_session_timeout; RESET lock_timeout').and_call_original
+      expect(connection).to receive(:execute).with("SET lock_timeout TO '15ms'").and_call_original
 
       subject.run { }
     end

spec/lib/gitlab/database/with_lock_retries_spec.rb

@@ -5,7 +5,7 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::WithLockRetries do
   let(:env) { {} }
   let(:logger) { Gitlab::Database::WithLockRetries::NULL_LOGGER }
-  let(:subject) { described_class.new(env: env, logger: logger, allow_savepoints: allow_savepoints, timing_configuration: timing_configuration) }
+  let(:subject) { described_class.new(connection: connection, env: env, logger: logger, allow_savepoints: allow_savepoints, timing_configuration: timing_configuration) }
   let(:allow_savepoints) { true }
   let(:connection) { ActiveRecord::Base.retrieve_connection }