1 files changed, 17 insertions, 94 deletions

diff --git a/lib/tasks/gitlab/db/lock_writes.rake b/lib/tasks/gitlab/db/lock_writes.rake
index 3a083036781..eb6d257cac5 100644
--- a/lib/tasks/gitlab/db/lock_writes.rake
+++ b/lib/tasks/gitlab/db/lock_writes.rake
@@ -2,22 +2,25 @@
 
 namespace :gitlab do
   namespace :db do
-    TRIGGER_FUNCTION_NAME = 'gitlab_schema_prevent_write'
-
     desc "GitLab | DB | Install prevent write triggers on all databases"
     task lock_writes: [:environment, 'gitlab:db:validate_config'] do
-      Gitlab::Database::EachDatabase.each_database_connection do |connection, database_name|
-        create_write_trigger_function(connection)
-
+      Gitlab::Database::EachDatabase.each_database_connection(include_shared: false) do |connection, database_name|
         schemas_for_connection = Gitlab::Database.gitlab_schemas_for_connection(connection)
         Gitlab::Database::GitlabSchema.tables_to_schema.each do |table_name, schema_name|
           # TODO: https://gitlab.com/gitlab-org/gitlab/-/issues/366834
           next if schema_name == :gitlab_geo
 
+          lock_writes_manager = Gitlab::Database::LockWritesManager.new(
+            table_name: table_name,
+            connection: connection,
+            database_name: database_name,
+            logger: Logger.new($stdout)
+          )
+
           if schemas_for_connection.include?(schema_name.to_sym)
-            drop_write_trigger(database_name, connection, table_name)
+            lock_writes_manager.unlock_writes
           else
-            create_write_trigger(database_name, connection, table_name)
+            lock_writes_manager.lock_writes
           end
         end
       end
@@ -30,96 +33,16 @@ namespace :gitlab do
           # TODO: https://gitlab.com/gitlab-org/gitlab/-/issues/366834
           next if schema_name == :gitlab_geo
 
-          drop_write_trigger(database_name, connection, table_name)
-        end
-        drop_write_trigger_function(connection)
-      end
-    end
-
-    def create_write_trigger_function(connection)
-      sql = <<-SQL
-        CREATE OR REPLACE FUNCTION #{TRIGGER_FUNCTION_NAME}()
-          RETURNS TRIGGER AS
-          $$
-          BEGIN
-            RAISE EXCEPTION 'Table: "%" is write protected within this Gitlab database.', TG_TABLE_NAME
-              USING ERRCODE = 'modifying_sql_data_not_permitted',
-              HINT = 'Make sure you are using the right database connection';
-          END
-          $$ LANGUAGE PLPGSQL
-      SQL
-
-      connection.execute(sql)
-    end
-
-    def drop_write_trigger_function(connection)
-      sql = <<-SQL
-        DROP FUNCTION IF EXISTS #{TRIGGER_FUNCTION_NAME}()
-      SQL
-
-      connection.execute(sql)
-    end
-
-    def create_write_trigger(database_name, connection, table_name)
-      puts "#{database_name}: '#{table_name}'... Lock Writes".color(:yellow)
-      sql = <<-SQL
-          DROP TRIGGER IF EXISTS #{write_trigger_name(table_name)} ON #{table_name};
-          CREATE TRIGGER #{write_trigger_name(table_name)}
-            BEFORE INSERT OR UPDATE OR DELETE OR TRUNCATE
-            ON #{table_name}
-            FOR EACH STATEMENT EXECUTE FUNCTION #{TRIGGER_FUNCTION_NAME}();
-      SQL
-
-      with_retries(connection) do
-        connection.execute(sql)
-      end
-    end
-
-    def drop_write_trigger(database_name, connection, table_name)
-      puts "#{database_name}: '#{table_name}'... Allow Writes".color(:green)
-      sql = <<-SQL
-        DROP TRIGGER IF EXISTS #{write_trigger_name(table_name)} ON #{table_name}
-      SQL
-
-      with_retries(connection) do
-        connection.execute(sql)
-      end
-    end
+          lock_writes_manager = Gitlab::Database::LockWritesManager.new(
+            table_name: table_name,
+            connection: connection,
+            database_name: database_name,
+            logger: Logger.new($stdout)
+          )
 
-    def with_retries(connection, &block)
-      with_statement_timeout_retries do
-        with_lock_retries(connection) do
-          yield
+          lock_writes_manager.unlock_writes
         end
       end
     end
-
-    def with_statement_timeout_retries(times = 5)
-      current_iteration = 1
-      begin
-        yield
-      rescue ActiveRecord::QueryCanceled => err
-        puts "Retrying after #{err.message}"
-
-        if current_iteration <= times
-          current_iteration += 1
-          retry
-        else
-          raise err
-        end
-      end
-    end
-
-    def with_lock_retries(connection, &block)
-      Gitlab::Database::WithLockRetries.new(
-        klass: "gitlab:db:lock_writes",
-        logger: Gitlab::AppLogger,
-        connection: connection
-      ).run(&block)
-    end
-
-    def write_trigger_name(table_name)
-      "gitlab_schema_write_trigger_for_#{table_name}"
-    end
   end
 end