Lock writes to detached partitions as well

This allows us to truncate such detached partitions
(truncate_legacy_tables)

As the detached partitions is in a different schema, we update the query
for table_locked_for_writes? to handle differnt schemas correctly. (PG
stores the schema separately in another column)

lib/gitlab/database/lock_writes_manager.rb

@@ -10,18 +10,34 @@ class LockWritesManager
       # See https://www.postgresql.org/message-id/16934.1568989957%40sss.pgh.pa.us
       EXPECTED_TRIGGER_RECORD_COUNT = 3
 
+      def self.tables_to_lock(connection)
+        Gitlab::Database::GitlabSchema.tables_to_schema.each do |table_name, schema_name|
+          yield table_name, schema_name
+        end
+
+        Gitlab::Database::SharedModel.using_connection(connection) do
+          Postgresql::DetachedPartition.find_each do |detached_partition|
+            yield detached_partition.fully_qualified_table_name, detached_partition.table_schema
+          end
+        end
+      end
+
       def initialize(table_name:, connection:, database_name:, logger: nil, dry_run: false)
         @table_name = table_name
         @connection = connection
         @database_name = database_name
         @logger = logger
         @dry_run = dry_run
+
+        @table_name_without_schema = ActiveRecord::ConnectionAdapters::PostgreSQL::Utils
+          .extract_schema_qualified_name(table_name)
+          .identifier
       end
 
       def table_locked_for_writes?(table_name)
         query = <<~SQL
             SELECT COUNT(*) from information_schema.triggers
-            WHERE event_object_table = '#{table_name}'
+            WHERE event_object_table = '#{table_name_without_schema}'
             AND trigger_name = '#{write_trigger_name(table_name)}'
         SQL
 
@@ -56,7 +72,7 @@ def unlock_writes
 
       private
 
-      attr_reader :table_name, :connection, :database_name, :logger, :dry_run
+      attr_reader :table_name, :connection, :database_name, :logger, :dry_run, :table_name_without_schema
 
       def execute_sql_statement(sql)
         if dry_run
@@ -99,7 +115,7 @@ def with_lock_retries(connection, &block)
       end
 
       def write_trigger_name(table_name)
-        "gitlab_schema_write_trigger_for_#{table_name}"
+        "gitlab_schema_write_trigger_for_#{table_name_without_schema}"
       end
     end
   end

lib/tasks/gitlab/db/lock_writes.rake

@@ -6,7 +6,8 @@ namespace :gitlab do
     task lock_writes: [:environment, 'gitlab:db:validate_config'] do
       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|
+
+        Gitlab::Database::LockWritesManager.tables_to_lock(connection) do |table_name, schema_name|
           # TODO: https://gitlab.com/gitlab-org/gitlab/-/issues/366834
           next if schema_name == :gitlab_geo
 
@@ -30,7 +31,7 @@ namespace :gitlab do
     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|
+        Gitlab::Database::LockWritesManager.tables_to_lock(connection) do |table_name, schema_name|
           # TODO: https://gitlab.com/gitlab-org/gitlab/-/issues/366834
           next if schema_name == :gitlab_geo
 

spec/lib/gitlab/database/lock_writes_manager_spec.rb

@@ -37,6 +37,14 @@
     it 'returns true for a table that is locked for writes' do
       expect { subject.lock_writes }.to change { subject.table_locked_for_writes?(test_table) }.from(false).to(true)
     end
+
+    context 'for detached partition tables in another schema' do
+      let(:test_table) { 'gitlab_partitions_dynamic._test_table_20220101' }
+
+      it 'returns true for a table that is locked for writes' do
+        expect { subject.lock_writes }.to change { subject.table_locked_for_writes?(test_table) }.from(false).to(true)
+      end
+    end
   end
 
   describe '#lock_writes' do

spec/tasks/gitlab/db/lock_writes_rake_spec.rb

@@ -19,6 +19,30 @@
   let(:main_connection) { ApplicationRecord.connection }
   let(:ci_connection) { Ci::ApplicationRecord.connection }
 
+  let(:detached_partition_table) { '_test_gitlab_main_part_20220101' }
+
+  before do
+    create_detached_partition_sql = <<~SQL
+      CREATE TABLE IF NOT EXISTS gitlab_partitions_dynamic._test_gitlab_main_part_20220101 (
+        id bigserial primary key not null
+      )
+    SQL
+
+    main_connection.execute(create_detached_partition_sql)
+    ci_connection.execute(create_detached_partition_sql)
+
+    Gitlab::Database::SharedModel.using_connection(main_connection) do
+      Postgresql::DetachedPartition.create!(
+        table_name: detached_partition_table,
+        drop_after: Time.current
+      )
+    end
+
+    allow(Gitlab::Database::GitlabSchema).to receive(:table_schema).and_call_original
+    allow(Gitlab::Database::GitlabSchema).to receive(:table_schema)
+      .with(detached_partition_table).and_return(:gitlab_main)
+  end
+
   context 'single database' do
     before do
       skip_if_multiple_databases_are_setup
@@ -46,6 +70,13 @@
   context 'multiple databases' do
     before do
       skip_if_multiple_databases_not_setup
+
+      Gitlab::Database::SharedModel.using_connection(ci_connection) do
+        Postgresql::DetachedPartition.create!(
+          table_name: detached_partition_table,
+          drop_after: Time.current
+        )
+      end
     end
 
     context 'when locking writes' do
@@ -87,6 +118,13 @@
           main_connection.execute("truncate ci_build_needs")
         end.to raise_error(ActiveRecord::StatementInvalid, /Table: "ci_build_needs" is write protected/)
       end
+
+      it 'prevents writes to detached partitions' do
+        run_rake_task('gitlab:db:lock_writes')
+        expect do
+          ci_connection.execute("INSERT INTO gitlab_partitions_dynamic.#{detached_partition_table} DEFAULT VALUES")
+        end.to raise_error(ActiveRecord::StatementInvalid, /Table: "#{detached_partition_table}" is write protected/)
+      end
     end
 
     context 'when running in dry_run mode' do
@@ -138,6 +176,14 @@
           main_connection.execute("delete from ci_builds")
         end.not_to raise_error
       end
+
+      it 'allows writes again to detached partitions' do
+        run_rake_task('gitlab:db:unlock_writes')
+
+        expect do
+          ci_connection.execute("INSERT INTO gitlab_partitions_dynamic._test_gitlab_main_part_20220101 DEFAULT VALUES")
+        end.not_to raise_error
+      end
     end
   end