Support multiple dbs when monitoring partitions

Update the process that monitors dynamic partitions to work with
multiple databases.

.rubocop_manual_todo.yml

@@ -2445,11 +2445,7 @@ Database/MultipleDatabases:
     - 'lib/gitlab/database/migrations/observers/query_log.rb'
     - 'lib/gitlab/database/multi_threaded_migration.rb'
     - 'lib/gitlab/database/partitioning_migration_helpers/backfill_partitioned_table.rb'
-    - 'lib/gitlab/database/partitioning/monthly_strategy.rb'
-    - 'lib/gitlab/database/partitioning/partition_manager.rb'
-    - 'lib/gitlab/database/partitioning/partition_creator.rb'
     - 'lib/gitlab/database/partitioning/replace_table.rb'
-    - 'lib/gitlab/database/partitioning/time_partition.rb'
     - 'lib/gitlab/database/postgres_hll/batch_distinct_counter.rb'
     - 'lib/gitlab/database/postgresql_adapter/dump_schema_versions_mixin.rb'
     - 'lib/gitlab/database/postgresql_database_tasks/load_schema_versions_mixin.rb'

app/workers/database/drop_detached_partitions_worker.rb

@@ -12,7 +12,7 @@ module Database
     def perform
       Gitlab::Database::Partitioning.drop_detached_partitions
     ensure
-      Gitlab::Database::Partitioning::PartitionMonitoring.new.report_metrics
+      Gitlab::Database::Partitioning.report_metrics
     end
   end
 end

app/workers/database/partition_management_worker.rb

@@ -14,7 +14,7 @@ module Database
     def perform
       Gitlab::Database::Partitioning.sync_partitions
     ensure
-      Gitlab::Database::Partitioning::PartitionMonitoring.new.report_metrics
+      Gitlab::Database::Partitioning.report_metrics
     end
   end
 end

lib/gitlab/database/each_database.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module EachDatabase
+      class << self
+        def each_database_connection
+          Gitlab::Database.databases.each_pair do |connection_name, connection_wrapper|
+            connection = connection_wrapper.scope.connection
+
+            with_shared_connection(connection, connection_name) do
+              yield connection, connection_name
+            end
+          end
+        end
+
+        def each_model_connection(models)
+          models.each do |model|
+            connection_name = model.connection.pool.db_config.name
+
+            with_shared_connection(model.connection, connection_name) do
+              yield model, connection_name
+            end
+          end
+        end
+
+        private
+
+        def with_shared_connection(connection, connection_name)
+          Gitlab::Database::SharedModel.using_connection(connection) do
+            Gitlab::AppLogger.debug(message: 'Switched database connection', connection_name: connection_name)
+
+            yield
+          end
+        end
+      end
+    end
+  end
+end

lib/gitlab/database/partitioning.rb

@@ -12,11 +12,31 @@ module Gitlab
       end
 
       def self.sync_partitions(models_to_sync = registered_models)
-        MultiDatabasePartitionManager.new(models_to_sync).sync_partitions
+        Gitlab::AppLogger.info(message: 'Syncing dynamic postgres partitions')
+
+        Gitlab::Database::EachDatabase.each_model_connection(models_to_sync) do |model|
+          PartitionManager.new(model).sync_partitions
+        end
+
+        Gitlab::AppLogger.info(message: 'Finished sync of dynamic postgres partitions')
+      end
+
+      def self.report_metrics(models_to_monitor = registered_models)
+        partition_monitoring = PartitionMonitoring.new
+
+        Gitlab::Database::EachDatabase.each_model_connection(models_to_monitor) do |model|
+          partition_monitoring.report_metrics_for_model(model)
+        end
       end
 
       def self.drop_detached_partitions
-        MultiDatabasePartitionDropper.new.drop_detached_partitions
+        Gitlab::AppLogger.info(message: 'Dropping detached postgres partitions')
+
+        Gitlab::Database::EachDatabase.each_database_connection do
+          DetachedPartitionDropper.new.perform
+        end
+
+        Gitlab::AppLogger.info(message: 'Finished dropping detached postgres partitions')
       end
     end
   end

lib/gitlab/database/partitioning/monthly_strategy.rb

@@ -96,10 +96,6 @@ module Gitlab
         def oldest_active_date
           (Date.today - retain_for).beginning_of_month
         end
-
-        def connection
-          ActiveRecord::Base.connection
-        end
       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/partitioning/multi_database_partition_manager.rb

-# frozen_string_literal: true
-
-module Gitlab
-  module Database
-    module Partitioning
-      class MultiDatabasePartitionManager
-        def initialize(models)
-          @models = models
-        end
-
-        def sync_partitions
-          Gitlab::AppLogger.info(message: "Syncing dynamic postgres partitions")
-
-          models.each do |model|
-            Gitlab::Database::SharedModel.using_connection(model.connection) do
-              Gitlab::AppLogger.debug(message: "Switched database connection",
-                                      connection_name: connection_name,
-                                      table_name: model.table_name)
-
-              PartitionManager.new(model).sync_partitions
-            end
-          end
-
-          Gitlab::AppLogger.info(message: "Finished sync of dynamic postgres partitions")
-        end
-
-        private
-
-        attr_reader :models
-
-        def connection_name
-          Gitlab::Database::SharedModel.connection.pool.db_config.name
-        end
-      end
-    end
-  end
-end

lib/gitlab/database/partitioning/partition_monitoring.rb

@@ -4,20 +4,12 @@ module Gitlab
   module Database
     module Partitioning
       class PartitionMonitoring
-        attr_reader :models
+        def report_metrics_for_model(model)
+          strategy = model.partitioning_strategy
 
-        def initialize(models = Gitlab::Database::Partitioning.registered_models)
-          @models = models
-        end
-
-        def report_metrics
-          models.each do |model|
-            strategy = model.partitioning_strategy
-
-            gauge_present.set({ table: model.table_name }, strategy.current_partitions.size)
-            gauge_missing.set({ table: model.table_name }, strategy.missing_partitions.size)
-            gauge_extra.set({ table: model.table_name }, strategy.extra_partitions.size)
-          end
+          gauge_present.set({ table: model.table_name }, strategy.current_partitions.size)
+          gauge_missing.set({ table: model.table_name }, strategy.missing_partitions.size)
+          gauge_extra.set({ table: model.table_name }, strategy.extra_partitions.size)
         end
 
         private

lib/gitlab/database/partitioning/time_partition.rb

@@ -87,7 +87,7 @@ module Gitlab
         end
 
         def conn
-          @conn ||= ActiveRecord::Base.connection
+          @conn ||= Gitlab::Database::SharedModel.connection
         end
       end
     end

spec/lib/gitlab/database/each_database_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::EachDatabase do
+  describe '.each_database' do
+    let(:expected_connections) do
+      Gitlab::Database.databases.map { |name, wrapper| [wrapper.scope.connection, name] }
+    end
+
+    it 'yields each connection after connecting SharedModel' do
+      expected_connections.each do |connection, _|
+        expect(Gitlab::Database::SharedModel).to receive(:using_connection).with(connection).and_yield
+      end
+
+      yielded_connections = []
+
+      described_class.each_database_connection do |connection, name|
+        yielded_connections << [connection, name]
+      end
+
+      expect(yielded_connections).to match_array(expected_connections)
+    end
+  end
+
+  describe '.each_model_connection' do
+    let(:model1) { double(connection: double, table_name: 'table1') }
+    let(:model2) { double(connection: double, table_name: 'table2') }
+
+    before do
+      allow(model1.connection).to receive_message_chain('pool.db_config.name').and_return('name1')
+      allow(model2.connection).to receive_message_chain('pool.db_config.name').and_return('name2')
+    end
+
+    it 'yields each model after connecting SharedModel' do
+      expect(Gitlab::Database::SharedModel).to receive(:using_connection).with(model1.connection).and_yield
+      expect(Gitlab::Database::SharedModel).to receive(:using_connection).with(model2.connection).and_yield
+
+      yielded_models = []
+
+      described_class.each_model_connection([model1, model2]) do |model, name|
+        yielded_models << [model, name]
+      end
+
+      expect(yielded_models).to match_array([[model1, 'name1'], [model2, 'name2']])
+    end
+  end
+end

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/multi_database_partition_manager_spec.rb

-# frozen_string_literal: true
-
-require 'spec_helper'
-
-RSpec.describe Gitlab::Database::Partitioning::MultiDatabasePartitionManager, '#sync_partitions' do
-  subject(:sync_partitions) { manager.sync_partitions }
-
-  let(:manager) { described_class.new(models) }
-  let(:models) { [model1, model2] }
-
-  let(:model1) { double('model1', connection: connection1, table_name: 'table1') }
-  let(:model2) { double('model2', connection: connection1, table_name: 'table2') }
-
-  let(:connection1) { double('connection1') }
-  let(:connection2) { double('connection2') }
-
-  let(:target_manager_class) { Gitlab::Database::Partitioning::PartitionManager }
-  let(:target_manager1) { double('partition manager') }
-  let(:target_manager2) { double('partition manager') }
-
-  before do
-    allow(manager).to receive(:connection_name).and_return('name')
-  end
-
-  it 'syncs model partitions, setting up the appropriate connection for each', :aggregate_failures do
-    expect(Gitlab::Database::SharedModel).to receive(:using_connection).with(model1.connection).and_yield.ordered
-    expect(target_manager_class).to receive(:new).with(model1).and_return(target_manager1).ordered
-    expect(target_manager1).to receive(:sync_partitions)
-
-    expect(Gitlab::Database::SharedModel).to receive(:using_connection).with(model2.connection).and_yield.ordered
-    expect(target_manager_class).to receive(:new).with(model2).and_return(target_manager2).ordered
-    expect(target_manager2).to receive(:sync_partitions)
-
-    sync_partitions
-  end
-end

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

@@ -4,9 +4,8 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Partitioning::PartitionMonitoring do
   describe '#report_metrics' do
-    subject { described_class.new(models).report_metrics }
+    subject { described_class.new.report_metrics_for_model(model) }
 
-    let(:models) { [model] }
     let(:model) { double(partitioning_strategy: partitioning_strategy, table_name: table) }
     let(:partitioning_strategy) { double(missing_partitions: missing_partitions, current_partitions: current_partitions, extra_partitions: extra_partitions) }
     let(:table) { "some_table" }

spec/lib/gitlab/database/partitioning_spec.rb

@@ -3,46 +3,121 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Partitioning do
+  include Database::PartitioningHelpers
+  include Database::TableSchemaHelpers
+
+  let(:connection) { ApplicationRecord.connection }
+
   describe '.sync_partitions' do
-    let(:partition_manager_class) { described_class::MultiDatabasePartitionManager }
-    let(:partition_manager) { double('partition manager') }
+    let(:table_names) { %w[partitioning_test1 partitioning_test2] }
+    let(:models) do
+      table_names.map do |table_name|
+        Class.new(ApplicationRecord) do
+          include PartitionedTable
+
+          self.table_name = table_name
+          partitioned_by :created_at, strategy: :monthly
+        end
+      end
+    end
+
+    before do
+      table_names.each do |table_name|
+        connection.execute(<<~SQL)
+          CREATE TABLE #{table_name} (
+            id serial not null,
+            created_at timestamptz not null,
+            PRIMARY KEY (id, created_at))
+          PARTITION BY RANGE (created_at);
+        SQL
+      end
+    end
+
+    it 'manages partitions for each given model' do
+      expect { described_class.sync_partitions(models)}
+        .to change { find_partitions(table_names.first).size }.from(0)
+        .and change { find_partitions(table_names.last).size }.from(0)
+    end
 
     context 'when no partitioned models are given' do
-      it 'calls the partition manager with the registered models' do
-        expect(partition_manager_class).to receive(:new)
+      let(:partition_manager_class) { described_class::PartitionManager }
+      let(:partition_manager) { double('partition manager') }
+      let(:model) { double('model') }
+
+      it 'manages partitions for each registered model' do
+        expect(Gitlab::Database::EachDatabase).to receive(:each_model_connection)
           .with(described_class.registered_models)
-          .and_return(partition_manager)
+          .and_yield(model)
 
+        expect(partition_manager_class).to receive(:new).with(model).and_return(partition_manager)
         expect(partition_manager).to receive(:sync_partitions)
 
         described_class.sync_partitions
       end
     end
+  end
 
-    context 'when partitioned models are given' do
-      it 'calls the partition manager with the given models' do
-        models = ['my special model']
+  describe '.report_metrics' do
+    let(:model1) { double('model') }
+    let(:model2) { double('model') }
 
-        expect(partition_manager_class).to receive(:new)
-          .with(models)
-          .and_return(partition_manager)
+    let(:partition_monitoring_class) { described_class::PartitionMonitoring }
 
-        expect(partition_manager).to receive(:sync_partitions)
+    context 'when no partitioned models are given' do
+      it 'reports metrics for each registered model' do
+        expect_next_instance_of(partition_monitoring_class) do |partition_monitor|
+          expect(partition_monitor).to receive(:report_metrics_for_model).with(model1)
+          expect(partition_monitor).to receive(:report_metrics_for_model).with(model2)
+        end
+
+        expect(Gitlab::Database::EachDatabase).to receive(:each_model_connection)
+          .with(described_class.registered_models)
+          .and_yield(model1)
+          .and_yield(model2)
+
+        described_class.report_metrics
+      end
+    end
+
+    context 'when partitioned models are given' do
+      it 'reports metrics for each given model' do
+        expect_next_instance_of(partition_monitoring_class) do |partition_monitor|
+          expect(partition_monitor).to receive(:report_metrics_for_model).with(model1)
+          expect(partition_monitor).to receive(:report_metrics_for_model).with(model2)
+        end
+
+        expect(Gitlab::Database::EachDatabase).to receive(:each_model_connection)
+          .with([model1, model2])
+          .and_yield(model1)
+          .and_yield(model2)
 
-        described_class.sync_partitions(models)
+        described_class.report_metrics([model1, model2])
       end
     end
   end
 
   describe '.drop_detached_partitions' do
-    let(:partition_dropper_class) { described_class::MultiDatabasePartitionDropper }
+    let(:table_names) { %w[detached_test_partition1 detached_test_partition2] }
 
-    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)
+    before do
+      table_names.each do |table_name|
+        connection.create_table("#{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}.#{table_name}")
+
+        Postgresql::DetachedPartition.create!(table_name: table_name, drop_after: 1.year.ago)
       end
+    end
+
+    it 'drops detached partitions for each database' do
+      expect(Gitlab::Database::EachDatabase).to receive(:each_database_connection).and_yield
+
+      expect { described_class.drop_detached_partitions }
+        .to change { Postgresql::DetachedPartition.count }.from(2).to(0)
+        .and change { table_exists?(table_names.first) }.from(true).to(false)
+        .and change { table_exists?(table_names.last) }.from(true).to(false)
+    end
 
-      described_class.drop_detached_partitions
+    def table_exists?(table_name)
+      table_oid(table_name).present?
     end
   end
 

spec/workers/database/drop_detached_partitions_worker_spec.rb

@@ -6,21 +6,19 @@ RSpec.describe Database::DropDetachedPartitionsWorker do
   describe '#perform' do
     subject { described_class.new.perform }
 
-    let(:monitoring) { instance_double('PartitionMonitoring', report_metrics: nil) }
-
     before do
       allow(Gitlab::Database::Partitioning).to receive(:drop_detached_partitions)
-      allow(Gitlab::Database::Partitioning::PartitionMonitoring).to receive(:new).and_return(monitoring)
+      allow(Gitlab::Database::Partitioning).to receive(:report_metrics)
     end
 
-    it 'delegates to Partitioning.drop_detached_partitions' do
+    it 'drops detached partitions' do
       expect(Gitlab::Database::Partitioning).to receive(:drop_detached_partitions)
 
       subject
     end
 
     it 'reports partition metrics' do
-      expect(monitoring).to receive(:report_metrics)
+      expect(Gitlab::Database::Partitioning).to receive(:report_metrics)
 
       subject
     end

spec/workers/database/partition_management_worker_spec.rb

@@ -6,20 +6,19 @@ RSpec.describe Database::PartitionManagementWorker do
   describe '#perform' do
     subject { described_class.new.perform }
 
-    let(:monitoring) { instance_double('PartitionMonitoring', report_metrics: nil) }
-
     before do
-      allow(Gitlab::Database::Partitioning::PartitionMonitoring).to receive(:new).and_return(monitoring)
+      allow(Gitlab::Database::Partitioning).to receive(:sync_partitions)
+      allow(Gitlab::Database::Partitioning).to receive(:report_metrics)
     end
 
-    it 'delegates to Partitioning' do
+    it 'syncs partitions' do
       expect(Gitlab::Database::Partitioning).to receive(:sync_partitions)
 
       subject
     end
 
     it 'reports partition metrics' do
-      expect(monitoring).to receive(:report_metrics)
+      expect(Gitlab::Database::Partitioning).to receive(:report_metrics)
 
       subject
     end