1 files changed, 18 insertions, 12 deletions

diff --git a/lib/gitlab/database/lock_writes_manager.rb b/lib/gitlab/database/lock_writes_manager.rb
index e3ae2892668..2e08e1ffb42 100644
--- a/lib/gitlab/database/lock_writes_manager.rb
+++ b/lib/gitlab/database/lock_writes_manager.rb
@@ -22,37 +22,38 @@ module Gitlab
         end
       end
 
-      def initialize(table_name:, connection:, database_name:, logger: nil, dry_run: false)
+      def initialize(table_name:, connection:, database_name:, with_retries: true, logger: nil, dry_run: false)
         @table_name = table_name
         @connection = connection
         @database_name = database_name
         @logger = logger
         @dry_run = dry_run
+        @with_retries = with_retries
 
         @table_name_without_schema = ActiveRecord::ConnectionAdapters::PostgreSQL::Utils
           .extract_schema_qualified_name(table_name)
           .identifier
       end
 
-      def table_locked_for_writes?(table_name)
+      def table_locked_for_writes?
         query = <<~SQL
             SELECT COUNT(*) from information_schema.triggers
             WHERE event_object_table = '#{table_name_without_schema}'
-            AND trigger_name = '#{write_trigger_name(table_name)}'
+            AND trigger_name = '#{write_trigger_name}'
         SQL
 
         connection.select_value(query) == EXPECTED_TRIGGER_RECORD_COUNT
       end
 
       def lock_writes
-        if table_locked_for_writes?(table_name)
+        if table_locked_for_writes?
           logger&.info "Skipping lock_writes, because #{table_name} is already locked for writes"
           return
         end
 
         logger&.info "Database: '#{database_name}', Table: '#{table_name}': Lock Writes".color(:yellow)
         sql_statement = <<~SQL
-          CREATE TRIGGER #{write_trigger_name(table_name)}
+          CREATE TRIGGER #{write_trigger_name}
             BEFORE INSERT OR UPDATE OR DELETE OR TRUNCATE
             ON #{table_name}
             FOR EACH STATEMENT EXECUTE FUNCTION #{TRIGGER_FUNCTION_NAME}();
@@ -64,7 +65,7 @@ module Gitlab
       def unlock_writes
         logger&.info "Database: '#{database_name}', Table: '#{table_name}': Allow Writes".color(:green)
         sql_statement = <<~SQL
-          DROP TRIGGER IF EXISTS #{write_trigger_name(table_name)} ON #{table_name};
+          DROP TRIGGER IF EXISTS #{write_trigger_name} ON #{table_name};
         SQL
 
         execute_sql_statement(sql_statement)
@@ -72,19 +73,23 @@ module Gitlab
 
       private
 
-      attr_reader :table_name, :connection, :database_name, :logger, :dry_run, :table_name_without_schema
+      attr_reader :table_name, :connection, :database_name, :logger, :dry_run, :table_name_without_schema, :with_retries
 
       def execute_sql_statement(sql)
         if dry_run
           logger&.info sql
-        else
-          with_retries(connection) do
+        elsif with_retries
+          raise "Cannot call lock_retries_helper if a transaction is already open" if connection.transaction_open?
+
+          run_with_retries(connection) do
             connection.execute(sql)
           end
+        else
+          connection.execute(sql)
         end
       end
 
-      def with_retries(connection, &block)
+      def run_with_retries(connection, &block)
         with_statement_timeout_retries do
           with_lock_retries(connection) do
             yield
@@ -110,11 +115,12 @@ module Gitlab
         Gitlab::Database::WithLockRetries.new(
           klass: "gitlab:db:lock_writes",
           logger: logger || Gitlab::AppLogger,
-          connection: connection
+          connection: connection,
+          allow_savepoints: false # this causes the WithLockRetries to fail if sub-transaction has been detected.
         ).run(&block)
       end
 
-      def write_trigger_name(table_name)
+      def write_trigger_name
         "gitlab_schema_write_trigger_for_#{table_name_without_schema}"
       end
     end