Support multiple dbs when dropping partitions

Update the process that drops detached partitions to work with multiple
databases.

app/workers/database/drop_detached_partitions_worker.rb

@@ -10,7 +10,7 @@ module Database
     idempotent!
 
     def perform
-      Gitlab::Database::Partitioning::DetachedPartitionDropper.new.perform
+      Gitlab::Database::Partitioning.drop_detached_partitions
     ensure
       Gitlab::Database::Partitioning::PartitionMonitoring.new.report_metrics
     end

lib/gitlab/database/partitioning.rb

@@ -14,6 +14,10 @@ module Gitlab
       def self.sync_partitions(models_to_sync = registered_models)
         MultiDatabasePartitionManager.new(models_to_sync).sync_partitions
       end
+
+      def self.drop_detached_partitions
+        MultiDatabasePartitionDropper.new.drop_detached_partitions
+      end
     end
   end
 end

lib/gitlab/database/partitioning/detached_partition_dropper.rb

@@ -7,18 +7,15 @@ module Gitlab
           return unless Feature.enabled?(:drop_detached_partitions, default_enabled: :yaml)
 
           Gitlab::AppLogger.info(message: "Checking for previously detached partitions to drop")
+
           Postgresql::DetachedPartition.ready_to_drop.find_each do |detached_partition|
-            conn.transaction do
+            connection.transaction do
               # Another process may have already dropped the table and deleted this entry
               next unless (detached_partition = Postgresql::DetachedPartition.lock.find_by(id: detached_partition.id))
 
-              unless check_partition_detached?(detached_partition)
-                Gitlab::AppLogger.error(message: "Attempt to drop attached database partition", partition_name: detached_partition.table_name)
-                detached_partition.destroy!
-                next
-              end
+              drop_detached_partition(detached_partition.table_name)
 
-              drop_one(detached_partition)
+              detached_partition.destroy!
             end
           rescue StandardError => e
             Gitlab::AppLogger.error(message: "Failed to drop previously detached partition",
@@ -30,25 +27,30 @@ module Gitlab
 
         private
 
-        def drop_one(detached_partition)
-          conn.transaction do
-            conn.execute(<<~SQL)
-              DROP TABLE #{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}.#{conn.quote_table_name(detached_partition.table_name)}
-            SQL
+        def drop_detached_partition(partition_name)
+          partition_identifier = qualify_partition_name(partition_name)
 
-            detached_partition.destroy!
+          if partition_detached?(partition_identifier)
+            connection.drop_table(partition_identifier, if_exists: true)
+
+            Gitlab::AppLogger.info(message: "Dropped previously detached partition", partition_name: partition_name)
+          else
+            Gitlab::AppLogger.error(message: "Attempt to drop attached database partition", partition_name: partition_name)
           end
-          Gitlab::AppLogger.info(message: "Dropped previously detached partition", partition_name: detached_partition.table_name)
         end
 
-        def check_partition_detached?(detached_partition)
+        def qualify_partition_name(table_name)
+          "#{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}.#{table_name}"
+        end
+
+        def partition_detached?(partition_identifier)
           # PostgresPartition checks the pg_inherits view, so our partition will only show here if it's still attached
           # and thus should not be dropped
-          !PostgresPartition.for_identifier("#{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}.#{detached_partition.table_name}").exists?
+          !Gitlab::Database::PostgresPartition.for_identifier(partition_identifier).exists?
         end
 
-        def conn
-          @conn ||= ApplicationRecord.connection
+        def connection
+          Postgresql::DetachedPartition.connection
         end
       end
     end

lib/gitlab/database/partitioning/multi_database_partition_dropper.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Partitioning
+      class MultiDatabasePartitionDropper
+        def drop_detached_partitions
+          Gitlab::AppLogger.info(message: "Dropping detached postgres partitions")
+
+          each_database_connection do |name, connection|
+            Gitlab::Database::SharedModel.using_connection(connection) do
+              Gitlab::AppLogger.debug(message: "Switched database connection", connection_name: name)
+
+              DetachedPartitionDropper.new.perform
+            end
+          end
+
+          Gitlab::AppLogger.info(message: "Finished dropping detached postgres partitions")
+        end
+
+        private
+
+        def each_database_connection
+          databases.each_pair do |name, connection_wrapper|
+            yield name, connection_wrapper.scope.connection
+          end
+        end
+
+        def databases
+          Gitlab::Database.databases
+        end
+      end
+    end
+  end
+end

lib/gitlab/database/shared_model.rb

@@ -2,6 +2,7 @@
 
 module Gitlab
   module Database
+    # This abstract class is used for models which need to exist in multiple de-composed databases.
     class SharedModel < ActiveRecord::Base
       self.abstract_class = true
 

spec/lib/gitlab/database/partitioning/detached_partition_dropper_spec.rb

@@ -84,6 +84,7 @@ RSpec.describe Gitlab::Database::Partitioning::DetachedPartitionDropper do
         before do
           stub_feature_flags(drop_detached_partitions: false)
         end
+
         it 'does not drop the partition' do
           subject.perform
 
@@ -162,8 +163,8 @@ RSpec.describe Gitlab::Database::Partitioning::DetachedPartitionDropper do
 
       context 'when the first drop returns an error' do
         it 'still drops the second partition' do
-          expect(subject).to receive(:drop_one).ordered.and_raise('injected error')
-          expect(subject).to receive(:drop_one).ordered.and_call_original
+          expect(subject).to receive(:drop_detached_partition).ordered.and_raise('injected error')
+          expect(subject).to receive(:drop_detached_partition).ordered.and_call_original
 
           subject.perform
 

spec/lib/gitlab/database/partitioning/multi_database_partition_dropper_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Partitioning::MultiDatabasePartitionDropper, '#drop_detached_partitions' do
+  subject(:drop_detached_partitions) { multi_db_dropper.drop_detached_partitions }
+
+  let(:multi_db_dropper) { described_class.new }
+
+  let(:connection_wrapper1) { double(scope: scope1) }
+  let(:connection_wrapper2) { double(scope: scope2) }
+
+  let(:scope1) { double(connection: connection1) }
+  let(:scope2) { double(connection: connection2) }
+
+  let(:connection1) { double('connection') }
+  let(:connection2) { double('connection') }
+
+  let(:dropper_class) { Gitlab::Database::Partitioning::DetachedPartitionDropper }
+  let(:dropper1) { double('partition dropper') }
+  let(:dropper2) { double('partition dropper') }
+
+  before do
+    allow(multi_db_dropper).to receive(:databases).and_return({ db1: connection_wrapper1, db2: connection_wrapper2 })
+  end
+
+  it 'drops detached partitions for each database' do
+    expect(Gitlab::Database::SharedModel).to receive(:using_connection).with(connection1).and_yield.ordered
+    expect(dropper_class).to receive(:new).and_return(dropper1).ordered
+    expect(dropper1).to receive(:perform)
+
+    expect(Gitlab::Database::SharedModel).to receive(:using_connection).with(connection2).and_yield.ordered
+    expect(dropper_class).to receive(:new).and_return(dropper2).ordered
+    expect(dropper2).to receive(:perform)
+
+    drop_detached_partitions
+  end
+end

spec/lib/gitlab/database/partitioning/partition_manager_spec.rb

@@ -176,7 +176,7 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
     end
 
     it 'detaches exactly one partition' do
-      expect { subject }.to change { find_partitions(my_model.table_name, schema: Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA).size }.from(9).to(8)
+      expect { subject }.to change { find_partitions(my_model.table_name).size }.from(9).to(8)
     end
 
     it 'detaches the old partition' do

spec/lib/gitlab/database/partitioning_spec.rb

@@ -33,4 +33,16 @@ RSpec.describe Gitlab::Database::Partitioning do
       end
     end
   end
+
+  describe '.drop_detached_partitions' do
+    let(:partition_dropper_class) { described_class::MultiDatabasePartitionDropper }
+
+    it 'delegates to the partition dropper' do
+      expect_next_instance_of(partition_dropper_class) do |partition_dropper|
+        expect(partition_dropper).to receive(:drop_detached_partitions)
+      end
+
+      described_class.drop_detached_partitions
+    end
+  end
 end

spec/workers/database/drop_detached_partitions_worker_spec.rb

@@ -6,16 +6,15 @@ RSpec.describe Database::DropDetachedPartitionsWorker do
   describe '#perform' do
     subject { described_class.new.perform }
 
-    let(:dropper) {  instance_double('DropDetachedPartitions', perform: nil) }
     let(:monitoring) { instance_double('PartitionMonitoring', report_metrics: nil) }
 
     before do
-      allow(Gitlab::Database::Partitioning::DetachedPartitionDropper).to receive(:new).and_return(dropper)
+      allow(Gitlab::Database::Partitioning).to receive(:drop_detached_partitions)
       allow(Gitlab::Database::Partitioning::PartitionMonitoring).to receive(:new).and_return(monitoring)
     end
 
-    it 'delegates to DropPartitionsPendingDrop' do
-      expect(dropper).to receive(:perform)
+    it 'delegates to Partitioning.drop_detached_partitions' do
+      expect(Gitlab::Database::Partitioning).to receive(:drop_detached_partitions)
 
       subject
     end