Merge branch 'ab/with-lock-retries-subtrans' into 'master'

Disallow subtransactions in with_lock_retries migration helper

See merge request gitlab-org/gitlab!68942

db/migrate/20210531053916_rename_instance_statistics_measurements.rb

 # frozen_string_literal: true
 
-class RenameInstanceStatisticsMeasurements < ActiveRecord::Migration[6.0]
-  include Gitlab::Database::MigrationHelpers
+class RenameInstanceStatisticsMeasurements < Gitlab::Database::Migration[1.0]
+  enable_lock_retries!
 
   def up
     rename_table_safely(:analytics_instance_statistics_measurements, :analytics_usage_trends_measurements)

db/migrate/20210621043337_rename_services_to_integrations.rb

 # frozen_string_literal: true
 
-class RenameServicesToIntegrations < ActiveRecord::Migration[6.1]
-  include Gitlab::Database::MigrationHelpers
+class RenameServicesToIntegrations < Gitlab::Database::Migration[1.0]
   include Gitlab::Database::SchemaHelpers
 
+  enable_lock_retries!
+
   # Function and trigger names match those migrated in:
   # - https://gitlab.com/gitlab-org/gitlab/-/merge_requests/49916
   # - https://gitlab.com/gitlab-org/gitlab/-/merge_requests/51852

db/migrate/20210902171406_add_latest_column_into_the_security_scans_table.rb

 # frozen_string_literal: true
 
 class AddLatestColumnIntoTheSecurityScansTable < Gitlab::Database::Migration[1.0]
+  enable_lock_retries!
+
   def up
-    with_lock_retries do
     add_column :security_scans, :latest, :boolean, default: true, null: false
   end
-  end
 
   def down
-    with_lock_retries do
     remove_column :security_scans, :latest
   end
-  end
 end

doc/development/database/rename_database_tables.md

@@ -60,6 +60,8 @@ Consider the next release as "Release N.M".
 Execute a standard migration (not a post-migration):
 
 ```ruby
+  enable_lock_retries!
+
   def up
     rename_table_safely(:issues, :tickets)
   end

lib/gitlab/database/migration_helpers.rb

@@ -380,6 +380,8 @@ module Gitlab
       # The timings can be controlled via the +timing_configuration+ parameter.
       # If the lock was not acquired within the retry period, a last attempt is made without using +lock_timeout+.
       #
+      # Note this helper uses subtransactions when run inside an already open transaction.
+      #
       # ==== Examples
       #   # Invoking without parameters
       #   with_lock_retries do
@@ -411,7 +413,8 @@ module Gitlab
         raise_on_exhaustion = !!kwargs.delete(:raise_on_exhaustion)
         merged_args = {
           klass: self.class,
-          logger: Gitlab::BackgroundMigration::Logger
+          logger: Gitlab::BackgroundMigration::Logger,
+          allow_savepoints: true
         }.merge(kwargs)
 
         Gitlab::Database::WithLockRetries.new(**merged_args)
@@ -1376,13 +1379,11 @@ into similar problems in the future (e.g. when new tables are created).
       # validate - Whether to validate the constraint in this call
       #
       def add_check_constraint(table, check, constraint_name, validate: true)
-        validate_check_constraint_name!(constraint_name)
-
         # Transactions would result in ALTER TABLE locks being held for the
         # duration of the transaction, defeating the purpose of this method.
-        if transaction_open?
-          raise 'add_check_constraint can not be run inside a transaction'
-        end
+        validate_not_in_transaction!(:add_check_constraint)
+
+        validate_check_constraint_name!(constraint_name)
 
         if check_constraint_exists?(table, constraint_name)
           warning_message = <<~MESSAGE
@@ -1427,6 +1428,10 @@ into similar problems in the future (e.g. when new tables are created).
       end
 
       def remove_check_constraint(table, constraint_name)
+        # This is technically not necessary, but aligned with add_check_constraint
+        # and allows us to continue use with_lock_retries here
+        validate_not_in_transaction!(:remove_check_constraint)
+
         validate_check_constraint_name!(constraint_name)
 
         # DROP CONSTRAINT requires an EXCLUSIVE lock

lib/gitlab/database/migration_helpers/v2.rb

@@ -6,6 +6,44 @@ module Gitlab
       module V2
         include Gitlab::Database::MigrationHelpers
 
+        # Executes the block with a retry mechanism that alters the +lock_timeout+ and +sleep_time+ between attempts.
+        # The timings can be controlled via the +timing_configuration+ parameter.
+        # If the lock was not acquired within the retry period, a last attempt is made without using +lock_timeout+.
+        #
+        # In order to retry the block, the method wraps the block into a transaction.
+        # Note it cannot be used inside an already open transaction and will raise an error in that case.
+        #
+        # ==== Examples
+        #   # Invoking without parameters
+        #   with_lock_retries do
+        #     drop_table :my_table
+        #   end
+        #
+        #   # Invoking with custom +timing_configuration+
+        #   t = [
+        #     [1.second, 1.second],
+        #     [2.seconds, 2.seconds]
+        #   ]
+        #
+        #   with_lock_retries(timing_configuration: t) do
+        #     drop_table :my_table # this will be retried twice
+        #   end
+        #
+        #   # Disabling the retries using an environment variable
+        #   > export DISABLE_LOCK_RETRIES=true
+        #
+        #   with_lock_retries do
+        #     drop_table :my_table # one invocation, it will not retry at all
+        #   end
+        #
+        # ==== Parameters
+        # * +timing_configuration+ - [[ActiveSupport::Duration, ActiveSupport::Duration], ...] lock timeout for the block, sleep time before the next iteration, defaults to `Gitlab::Database::WithLockRetries::DEFAULT_TIMING_CONFIGURATION`
+        # * +logger+ - [Gitlab::JsonLogger]
+        # * +env+ - [Hash] custom environment hash, see the example with `DISABLE_LOCK_RETRIES`
+        def with_lock_retries(*args, **kwargs, &block)
+          super(*args, **kwargs.merge(allow_savepoints: false), &block)
+        end
+
         # Renames a column without requiring downtime.
         #
         # Concurrent renames work by using database triggers to ensure both the

lib/gitlab/database/partitioning_migration_helpers/foreign_key_helpers.rb

@@ -6,6 +6,8 @@ module Gitlab
       module ForeignKeyHelpers
         include ::Gitlab::Database::SchemaHelpers
 
+        ERROR_SCOPE = 'foreign keys'
+
         # Adds a foreign key with only minimal locking on the tables involved.
         #
         # In concept it works similarly to add_concurrent_foreign_key, but we have
@@ -32,6 +34,8 @@ module Gitlab
         # name - The name of the foreign key.
         #
         def add_concurrent_partitioned_foreign_key(source, target, column:, on_delete: :cascade, name: nil)
+          assert_not_in_transaction_block(scope: ERROR_SCOPE)
+
           partition_options = {
             column: column,
             on_delete: on_delete,

lib/gitlab/database/partitioning_migration_helpers/index_helpers.rb

@@ -7,6 +7,8 @@ module Gitlab
         include Gitlab::Database::MigrationHelpers
         include Gitlab::Database::SchemaHelpers
 
+        ERROR_SCOPE = 'index'
+
         # Concurrently creates a new index on a partitioned table. In concept this works similarly to
         # `add_concurrent_index`, and won't block reads or writes on the table while the index is being built.
         #
@@ -21,6 +23,8 @@ module Gitlab
         #
         # See Rails' `add_index` for more info on the available arguments.
         def add_concurrent_partitioned_index(table_name, column_names, options = {})
+          assert_not_in_transaction_block(scope: ERROR_SCOPE)
+
           raise ArgumentError, 'A name is required for indexes added to partitioned tables' unless options[:name]
 
           partitioned_table = find_partitioned_table(table_name)
@@ -57,6 +61,8 @@ module Gitlab
         #
         #     remove_concurrent_partitioned_index_by_name :users, 'index_name_goes_here'
         def remove_concurrent_partitioned_index_by_name(table_name, index_name)
+          assert_not_in_transaction_block(scope: ERROR_SCOPE)
+
           find_partitioned_table(table_name)
 
           unless index_name_exists?(table_name, index_name)

lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb

@@ -431,7 +431,7 @@ module Gitlab
           replace_table = Gitlab::Database::Partitioning::ReplaceTable.new(original_table_name.to_s,
               replacement_table_name, replaced_table_name, primary_key_name)
 
-          with_lock_retries do
+          transaction do
             drop_sync_trigger(original_table_name)
 
             replace_table.perform do |sql|

lib/gitlab/database/rename_table_helpers.rb

@@ -4,27 +4,27 @@ module Gitlab
   module Database
     module RenameTableHelpers
       def rename_table_safely(old_table_name, new_table_name)
-        with_lock_retries do
+        transaction do
           rename_table(old_table_name, new_table_name)
           execute("CREATE VIEW #{old_table_name} AS SELECT * FROM #{new_table_name}")
         end
       end
 
       def undo_rename_table_safely(old_table_name, new_table_name)
-        with_lock_retries do
+        transaction do
           execute("DROP VIEW IF EXISTS #{old_table_name}")
           rename_table(new_table_name, old_table_name)
         end
       end
 
       def finalize_table_rename(old_table_name, new_table_name)
-        with_lock_retries do
+        transaction do
           execute("DROP VIEW IF EXISTS #{old_table_name}")
         end
       end
 
       def undo_finalize_table_rename(old_table_name, new_table_name)
-        with_lock_retries do
+        transaction do
           execute("CREATE VIEW #{old_table_name} AS SELECT * FROM #{new_table_name}")
         end
       end

lib/gitlab/database/with_lock_retries.rb

@@ -61,9 +61,10 @@ module Gitlab
         [10.seconds, 10.minutes]
       ].freeze
 
-      def initialize(logger: NULL_LOGGER, timing_configuration: DEFAULT_TIMING_CONFIGURATION, klass: nil, env: ENV)
+      def initialize(logger: NULL_LOGGER, allow_savepoints: true, timing_configuration: DEFAULT_TIMING_CONFIGURATION, klass: nil, env: ENV)
         @logger = logger
         @klass = klass
+        @allow_savepoints = allow_savepoints
         @timing_configuration = timing_configuration
         @env = env
         @current_iteration = 1
@@ -122,6 +123,8 @@ module Gitlab
       end
 
       def run_block_with_lock_timeout
+        raise "WithLockRetries should not run inside already open transaction" if ActiveRecord::Base.connection.transaction_open? && @allow_savepoints.blank?
+
         ActiveRecord::Base.transaction(requires_new: true) do # rubocop:disable Performance/ActiveRecordSubtransactions
           execute("SET LOCAL lock_timeout TO '#{current_lock_timeout_in_ms}ms'")
 

spec/lib/gitlab/database/migration_helpers/v2_spec.rb

@@ -11,6 +11,8 @@ RSpec.describe Gitlab::Database::MigrationHelpers::V2 do
 
   before do
     allow(migration).to receive(:puts)
+
+    allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
   end
 
   shared_examples_for 'Setting up to rename a column' do
@@ -218,4 +220,49 @@ RSpec.describe Gitlab::Database::MigrationHelpers::V2 do
       let(:added_column) { :original }
     end
   end
+
+  describe '#with_lock_retries' do
+    let(:model) do
+      ActiveRecord::Migration.new.extend(described_class)
+    end
+
+    let(:buffer) { StringIO.new }
+    let(:in_memory_logger) { Gitlab::JsonLogger.new(buffer) }
+    let(:env) { { 'DISABLE_LOCK_RETRIES' => 'true' } }
+
+    it 'sets the migration class name in the logs' do
+      model.with_lock_retries(env: env, logger: in_memory_logger) { }
+
+      buffer.rewind
+      expect(buffer.read).to include("\"class\":\"#{model.class}\"")
+    end
+
+    where(raise_on_exhaustion: [true, false])
+
+    with_them do
+      it 'sets raise_on_exhaustion as requested' do
+        with_lock_retries = double
+        expect(Gitlab::Database::WithLockRetries).to receive(:new).and_return(with_lock_retries)
+        expect(with_lock_retries).to receive(:run).with(raise_on_exhaustion: raise_on_exhaustion)
+
+        model.with_lock_retries(env: env, logger: in_memory_logger, raise_on_exhaustion: raise_on_exhaustion) { }
+      end
+    end
+
+    it 'does not raise on exhaustion by default' do
+      with_lock_retries = double
+      expect(Gitlab::Database::WithLockRetries).to receive(:new).and_return(with_lock_retries)
+      expect(with_lock_retries).to receive(:run).with(raise_on_exhaustion: false)
+
+      model.with_lock_retries(env: env, logger: in_memory_logger) { }
+    end
+
+    it 'defaults to disallowing subtransactions' do
+      with_lock_retries = double
+      expect(Gitlab::Database::WithLockRetries).to receive(:new).with(hash_including(allow_savepoints: false)).and_return(with_lock_retries)
+      expect(with_lock_retries).to receive(:run).with(raise_on_exhaustion: false)
+
+      model.with_lock_retries(env: env, logger: in_memory_logger) { }
+    end
+  end
 end

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -2310,8 +2310,6 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
       expect(buffer.read).to include("\"class\":\"#{model.class}\"")
     end
 
-    using RSpec::Parameterized::TableSyntax
-
     where(raise_on_exhaustion: [true, false])
 
     with_them do
@@ -2331,6 +2329,15 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
 
       model.with_lock_retries(env: env, logger: in_memory_logger) { }
     end
+
+    it 'defaults to allowing subtransactions' do
+      with_lock_retries = double
+
+      expect(Gitlab::Database::WithLockRetries).to receive(:new).with(hash_including(allow_savepoints: true)).and_return(with_lock_retries)
+      expect(with_lock_retries).to receive(:run).with(raise_on_exhaustion: false)
+
+      model.with_lock_retries(env: env, logger: in_memory_logger) { }
+    end
   end
 
   describe '#backfill_iids' do
@@ -2683,6 +2690,10 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
   end
 
   describe '#remove_check_constraint' do
+    before do
+      allow(model).to receive(:transaction_open?).and_return(false)
+    end
+
     it 'removes the constraint' do
       drop_sql = /ALTER TABLE test_table\s+DROP CONSTRAINT IF EXISTS check_name/
 

spec/lib/gitlab/database/partitioning_migration_helpers/foreign_key_helpers_spec.rb

@@ -27,6 +27,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::ForeignKeyHelpers
 
   before do
     allow(migration).to receive(:puts)
+    allow(migration).to receive(:transaction_open?).and_return(false)
 
     connection.execute(<<~SQL)
       CREATE TABLE #{target_table_name} (
@@ -141,5 +142,15 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::ForeignKeyHelpers
           .with(source_table_name, target_table_name, options)
       end
     end
+
+    context 'when run inside a transaction block' do
+      it 'raises an error' do
+        expect(migration).to receive(:transaction_open?).and_return(true)
+
+        expect do
+          migration.add_concurrent_partitioned_foreign_key(source_table_name, target_table_name, column: column_name)
+        end.to raise_error(/can not be run inside a transaction/)
+      end
+    end
   end
 end

spec/lib/gitlab/database/partitioning_migration_helpers/index_helpers_spec.rb

@@ -20,6 +20,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::IndexHelpers do
 
   before do
     allow(migration).to receive(:puts)
+    allow(migration).to receive(:transaction_open?).and_return(false)
 
     connection.execute(<<~SQL)
       CREATE TABLE #{table_name} (
@@ -127,6 +128,16 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::IndexHelpers do
         end.to raise_error(ArgumentError, /#{table_name} is not a partitioned table/)
       end
     end
+
+    context 'when run inside a transaction block' do
+      it 'raises an error' do
+        expect(migration).to receive(:transaction_open?).and_return(true)
+
+        expect do
+          migration.add_concurrent_partitioned_index(table_name, column_name)
+        end.to raise_error(/can not be run inside a transaction/)
+      end
+    end
   end
 
   describe '#remove_concurrent_partitioned_index_by_name' do
@@ -182,5 +193,15 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::IndexHelpers do
         end.to raise_error(ArgumentError, /#{table_name} is not a partitioned table/)
       end
     end
+
+    context 'when run inside a transaction block' do
+      it 'raises an error' do
+        expect(migration).to receive(:transaction_open?).and_return(true)
+
+        expect do
+          migration.remove_concurrent_partitioned_index_by_name(table_name, index_name)
+        end.to raise_error(/can not be run inside a transaction/)
+      end
+    end
   end
 end

spec/lib/gitlab/database/with_lock_retries_spec.rb

@@ -5,7 +5,8 @@ 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, timing_configuration: timing_configuration) }
+  let(:subject) { described_class.new(env: env, logger: logger, allow_savepoints: allow_savepoints, timing_configuration: timing_configuration) }
+  let(:allow_savepoints) { true }
 
   let(:timing_configuration) do
     [
@@ -256,4 +257,20 @@ RSpec.describe Gitlab::Database::WithLockRetries do
       subject.run { }
     end
   end
+
+  context 'Stop using subtransactions - allow_savepoints: false' do
+    let(:allow_savepoints) { false }
+
+    it 'prevents running inside already open transaction' do
+      allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(true)
+
+      expect { subject.run { } }.to raise_error(/should not run inside already open transaction/)
+    end
+
+    it 'does not raise the error if not inside open transaction' do
+      allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
+
+      expect { subject.run { } }.not_to raise_error
+    end
+  end
 end