1 files changed, 29 insertions, 13 deletions

diff --git a/lib/gitlab/database/tables_truncate.rb b/lib/gitlab/database/tables_truncate.rb
index 8380bf23899..807ecdb862a 100644
--- a/lib/gitlab/database/tables_truncate.rb
+++ b/lib/gitlab/database/tables_truncate.rb
@@ -19,24 +19,32 @@ module Gitlab
 
         logger&.info "DRY RUN:" if dry_run
 
-        connection = Gitlab::Database.database_base_models[database_name].connection
-
         schemas_for_connection = Gitlab::Database.gitlab_schemas_for_connection(connection)
         tables_to_truncate = Gitlab::Database::GitlabSchema.tables_to_schema.reject do |_, schema_name|
-          (GITLAB_SCHEMAS_TO_IGNORE.union(schemas_for_connection)).include?(schema_name)
+          GITLAB_SCHEMAS_TO_IGNORE.union(schemas_for_connection).include?(schema_name)
         end.keys
 
+        Gitlab::Database::SharedModel.using_connection(connection) do
+          Postgresql::DetachedPartition.find_each do |detached_partition|
+            next if GITLAB_SCHEMAS_TO_IGNORE.union(schemas_for_connection).include?(detached_partition.table_schema)
+
+            tables_to_truncate << detached_partition.fully_qualified_table_name
+          end
+        end
+
         tables_sorted = Gitlab::Database::TablesSortedByForeignKeys.new(connection, tables_to_truncate).execute
         # Checking if all the tables have the write-lock triggers
         # to make sure we are deleting the right tables on the right database.
         tables_sorted.flatten.each do |table_name|
-          query = <<~SQL
-            SELECT COUNT(*) from information_schema.triggers
-            WHERE event_object_table = '#{table_name}'
-            AND trigger_name = 'gitlab_schema_write_trigger_for_#{table_name}'
-          SQL
-
-          if connection.select_value(query) == 0
+          lock_writes_manager = Gitlab::Database::LockWritesManager.new(
+            table_name: table_name,
+            connection: connection,
+            database_name: database_name,
+            logger: logger,
+            dry_run: dry_run
+          )
+
+          unless lock_writes_manager.table_locked_for_writes?(table_name)
             raise "Table '#{table_name}' is not locked for writes. Run the rake task gitlab:db:lock_writes first"
           end
         end
@@ -51,18 +59,26 @@ module Gitlab
         # min_batch_size is the minimum number of new tables to truncate at each stage.
         # But in each stage we have also have to truncate the already truncated tables in the previous stages
         logger&.info "Truncating legacy tables for the database #{database_name}"
-        truncate_tables_in_batches(connection, tables_sorted, min_batch_size)
+        truncate_tables_in_batches(tables_sorted)
       end
 
       private
 
       attr_accessor :database_name, :min_batch_size, :logger, :dry_run, :until_table
 
-      def truncate_tables_in_batches(connection, tables_sorted, min_batch_size)
+      def connection
+        @connection ||= Gitlab::Database.database_base_models[database_name].connection
+      end
+
+      def truncate_tables_in_batches(tables_sorted)
         truncated_tables = []
 
         tables_sorted.flatten.each do |table|
-          sql_statement = "SELECT set_config('lock_writes.#{table}', 'false', false)"
+          table_name_without_schema = ActiveRecord::ConnectionAdapters::PostgreSQL::Utils
+            .extract_schema_qualified_name(table)
+            .identifier
+
+          sql_statement = "SELECT set_config('lock_writes.#{table_name_without_schema}', 'false', false)"
           logger&.info(sql_statement)
           connection.execute(sql_statement) unless dry_run
         end