Fix partition manager to run against selected dbs

Update the partition management to allow it to be only run against a
subset of the available databases. This is needed for rake tasks where
we might only interact with `main` or `ci` at one time.

lib/gitlab/database/each_database.rb

@@ -17,14 +17,16 @@ module Gitlab
           end
         end
 
-        def each_model_connection(models, &blk)
+        def each_model_connection(models, only_on: nil, &blk)
+          selected_databases = Array.wrap(only_on).map(&:to_sym)
+
           models.each do |model|
             # If model is shared, iterate all available base connections
             # Example: `LooseForeignKeys::DeletedRecord`
             if model < ::Gitlab::Database::SharedModel
-              with_shared_model_connections(model, &blk)
+              with_shared_model_connections(model, selected_databases, &blk)
             else
-              with_model_connection(model, &blk)
+              with_model_connection(model, selected_databases, &blk)
             end
           end
         end
@@ -43,20 +45,24 @@ module Gitlab
           end
         end
 
-        def with_shared_model_connections(shared_model, &blk)
+        def with_shared_model_connections(shared_model, selected_databases, &blk)
           Gitlab::Database.database_base_models.each_pair do |connection_name, connection_model|
             if shared_model.limit_connection_names
               next unless shared_model.limit_connection_names.include?(connection_name.to_sym)
             end
 
+            next if selected_databases.present? && !selected_databases.include?(connection_name.to_sym)
+
             with_shared_connection(connection_model.connection, connection_name) do
               yield shared_model, connection_name
             end
           end
         end
 
-        def with_model_connection(model, &blk)
-          connection_name = model.connection.pool.db_config.name
+        def with_model_connection(model, selected_databases, &blk)
+          connection_name = model.connection_db_config.name
+
+          return if selected_databases.present? && !selected_databases.include?(connection_name.to_sym)
 
           with_shared_connection(model.connection, connection_name) do
             yield model, connection_name

lib/gitlab/database/partitioning.rb

@@ -26,10 +26,10 @@ module Gitlab
           # ignore - happens when Rake tasks yet have to create a database, e.g. for testing
         end
 
-        def sync_partitions(models_to_sync = registered_for_sync)
+        def sync_partitions(models_to_sync = registered_for_sync, only_on: nil)
           Gitlab::AppLogger.info(message: 'Syncing dynamic postgres partitions')
 
-          Gitlab::Database::EachDatabase.each_model_connection(models_to_sync) do |model|
+          Gitlab::Database::EachDatabase.each_model_connection(models_to_sync, only_on: only_on) do |model|
             PartitionManager.new(model).sync_partitions
           end
 

spec/lib/gitlab/database/each_database_spec.rb

@@ -105,8 +105,8 @@ RSpec.describe Gitlab::Database::EachDatabase do
         allow(main_model).to receive(:connection).and_return(main_connection)
         allow(ci_model).to receive(:connection).and_return(ci_connection)
 
-        allow(main_connection).to receive_message_chain('pool.db_config.name').and_return('main')
-        allow(ci_connection).to receive_message_chain('pool.db_config.name').and_return('ci')
+        allow(main_model).to receive_message_chain('connection_db_config.name').and_return('main')
+        allow(ci_model).to receive_message_chain('connection_db_config.name').and_return('ci')
       end
 
       it 'yields each model after connecting SharedModel' do
@@ -117,10 +117,44 @@ RSpec.describe Gitlab::Database::EachDatabase do
       end
     end
 
-    def expect_yielded_models(models_to_iterate, expected_values)
+    context 'when the database connections are limited by the only_on option' do
+      let(:shared_model) { Class.new(Gitlab::Database::SharedModel) }
+      let(:main_model) { Class.new(ActiveRecord::Base) }
+      let(:ci_model) { Class.new(Ci::ApplicationRecord) }
+
+      before do
+        allow(Gitlab::Database).to receive(:database_base_models)
+          .and_return({ main: ActiveRecord::Base, ci: Ci::ApplicationRecord }.with_indifferent_access)
+
+        allow(main_model).to receive_message_chain('connection_db_config.name').and_return('main')
+        allow(ci_model).to receive_message_chain('connection_db_config.name').and_return('ci')
+      end
+
+      context 'when a single name is passed in' do
+        it 'yields models only connected to the given database' do
+          expect_yielded_models([main_model, ci_model, shared_model], [
+            { model: ci_model, connection: Ci::ApplicationRecord.connection, name: 'ci' },
+            { model: shared_model, connection: Ci::ApplicationRecord.connection, name: 'ci' }
+          ], only_on: 'ci')
+        end
+      end
+
+      context 'when a list of names are passed in' do
+        it 'yields models only connected to the given databases' do
+          expect_yielded_models([main_model, ci_model, shared_model], [
+            { model: main_model, connection: ActiveRecord::Base.connection, name: 'main' },
+            { model: ci_model, connection: Ci::ApplicationRecord.connection, name: 'ci' },
+            { model: shared_model, connection: ActiveRecord::Base.connection, name: 'main' },
+            { model: shared_model, connection: Ci::ApplicationRecord.connection, name: 'ci' }
+          ], only_on: %i[main ci])
+        end
+      end
+    end
+
+    def expect_yielded_models(models_to_iterate, expected_values, only_on: nil)
       times_yielded = 0
 
-      described_class.each_model_connection(models_to_iterate) do |model, name|
+      described_class.each_model_connection(models_to_iterate, only_on: only_on) do |model, name|
         expected = expected_values[times_yielded]
 
         expect(model).to be(expected[:model])

spec/lib/gitlab/database/partitioning_spec.rb

@@ -109,6 +109,20 @@ RSpec.describe Gitlab::Database::Partitioning do
           .and change { find_partitions(table_names.last).size }.from(0)
       end
     end
+
+    context 'when only a specific database is requested' do
+      before do
+        allow(models.first).to receive_message_chain('connection_db_config.name').and_return('main')
+        allow(models.last).to receive_message_chain('connection_db_config.name').and_return('ci')
+      end
+
+      it 'manages partitions for models for the given database', :aggregate_failures do
+        expect { described_class.sync_partitions(models, only_on: 'ci') }
+          .to change { find_partitions(table_names.last).size }.from(0)
+
+        expect(find_partitions(table_names.first).size).to eq(0)
+      end
+    end
   end
 
   describe '.report_metrics' do