1 files changed, 119 insertions, 0 deletions
diff --git a/lib/tasks/gitlab/db/lock_writes.rake b/lib/tasks/gitlab/db/lock_writes.rake
new file mode 100644
index 00000000000..b57c2860fe3
--- /dev/null
+++ b/lib/tasks/gitlab/db/lock_writes.rake
@@ -0,0 +1,119 @@
+# frozen_string_literal: true
+
+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)
+
+        schemas_for_connection = Gitlab::Database.gitlab_schemas_for_connection(connection)
+        Gitlab::Database::GitlabSchema.tables_to_schema.each do |table_name, schema_name|
+          if schemas_for_connection.include?(schema_name.to_sym)
+            drop_write_trigger(database_name, connection, table_name)
+          else
+            create_write_trigger(database_name, connection, table_name)
+          end
+        end
+      end
+    end
+
+    desc "GitLab | DB | Remove all triggers that prevents writes from all databases"
+    task unlock_writes: :environment do
+      Gitlab::Database::EachDatabase.each_database_connection do |connection, database_name|
+        Gitlab::Database::GitlabSchema.tables_to_schema.each do |table_name, schema_name|
+          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
+
+    def with_retries(connection, &block)
+      with_statement_timeout_retries do
+        with_lock_retries(connection) do
+          yield
+        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