Support multiple databases for partitioning

Update PartitionManager to support partition management across multiple
databases.

app/models/concerns/partitioned_table.rb

@@ -14,8 +14,6 @@ module PartitionedTable
       strategy_class = PARTITIONING_STRATEGIES[strategy.to_sym] || raise(ArgumentError, "Unknown partitioning strategy: #{strategy}")
 
       @partitioning_strategy = strategy_class.new(self, partitioning_key, **kwargs)
-
-      Gitlab::Database::Partitioning::PartitionManager.register(self)
     end
   end
 end

app/workers/database/partition_management_worker.rb

@@ -12,7 +12,7 @@ module Database
     idempotent!
 
     def perform
-      Gitlab::Database::Partitioning::PartitionManager.new.sync_partitions
+      Gitlab::Database::Partitioning.sync_partitions
     ensure
       Gitlab::Database::Partitioning::PartitionMonitoring.new.report_metrics
     end

config/initializers/postgres_partitioning.rb

 # frozen_string_literal: true
 
-# Make sure we have loaded partitioned models here
-# (even with eager loading disabled).
-
-Gitlab::Database::Partitioning::PartitionManager.register(AuditEvent)
-Gitlab::Database::Partitioning::PartitionManager.register(WebHookLog)
-Gitlab::Database::Partitioning::PartitionManager.register(LooseForeignKeys::DeletedRecord)
-
-if Gitlab.ee?
-  Gitlab::Database::Partitioning::PartitionManager.register(IncidentManagement::PendingEscalations::Alert)
-  Gitlab::Database::Partitioning::PartitionManager.register(IncidentManagement::PendingEscalations::Issue)
-end
-
 begin
-  Gitlab::Database::Partitioning::PartitionManager.new.sync_partitions unless ENV['DISABLE_POSTGRES_PARTITION_CREATION_ON_STARTUP']
+  Gitlab::Database::Partitioning.sync_partitions unless ENV['DISABLE_POSTGRES_PARTITION_CREATION_ON_STARTUP']
 rescue ActiveRecord::ActiveRecordError, PG::Error
   # ignore - happens when Rake tasks yet have to create a database, e.g. for testing
 end

lib/gitlab/database/partitioning.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Partitioning
+      def self.sync_partitions(partitioned_models = default_partitioned_models)
+        MultiDatabasePartitionManager.new(partitioned_models).sync_partitions
+      end
+
+      def self.default_partitioned_models
+        @default_partitioned_models ||= core_partitioned_models.union(ee_partitioned_models)
+      end
+
+      def self.core_partitioned_models
+        @core_partitioned_models ||= Set[
+          ::AuditEvent,
+          ::WebHookLog
+        ].freeze
+      end
+
+      def self.ee_partitioned_models
+        return Set.new.freeze unless Gitlab.ee?
+
+        @ee_partitioned_models ||= Set[
+          ::IncidentManagement::PendingEscalations::Alert,
+          ::IncidentManagement::PendingEscalations::Issue
+        ].freeze
+      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
+          return if models.empty?
+
+          each_database_connection do
+            PartitionManager.new(models).sync_partitions
+          end
+        end
+
+        private
+
+        attr_reader :models
+
+        def each_database_connection(&block)
+          original_db_config = ActiveRecord::Base.connection_db_config # rubocop:disable Database/MultipleDatabases
+
+          begin
+            with_each_connection(&block)
+          ensure
+            ActiveRecord::Base.establish_connection(original_db_config) # rubocop:disable Database/MultipleDatabases
+          end
+        end
+
+        def with_each_connection
+          Gitlab::Database.db_config_names.each do |db_name|
+            config_for_db_name = ActiveRecord::Base.configurations.configs_for(env_name: Rails.env, name: db_name) # rubocop:disable Database/MultipleDatabases
+
+            ActiveRecord::Base.establish_connection(config_for_db_name)
+
+            yield
+          end
+        end
+      end
+    end
+  end
+end

lib/gitlab/database/partitioning/partition_manager.rb

@@ -6,23 +6,11 @@ module Gitlab
       class PartitionManager
         UnsafeToDetachPartitionError = Class.new(StandardError)
 
-        def self.register(model)
-          raise ArgumentError, "Only models with a #partitioning_strategy can be registered." unless model.respond_to?(:partitioning_strategy)
-
-          models << model
-        end
-
-        def self.models
-          @models ||= Set.new
-        end
-
         LEASE_TIMEOUT = 1.minute
         MANAGEMENT_LEASE_KEY = 'database_partition_management_%s'
         RETAIN_DETACHED_PARTITIONS_FOR = 1.week
 
-        attr_reader :models
-
-        def initialize(models = self.class.models)
+        def initialize(models)
           @models = models
         end
 
@@ -53,6 +41,8 @@ module Gitlab
 
         private
 
+        attr_reader :models
+
         def missing_partitions(model)
           return [] unless connection.table_exists?(model.table_name)
 

lib/gitlab/database/partitioning/partition_monitoring.rb

@@ -6,7 +6,7 @@ module Gitlab
       class PartitionMonitoring
         attr_reader :models
 
-        def initialize(models = PartitionManager.models)
+        def initialize(models = Gitlab::Database::Partitioning.default_partitioned_models)
           @models = models
         end
 

lib/tasks/gitlab/db.rake

@@ -118,7 +118,7 @@ namespace :gitlab do
 
     desc 'Create missing dynamic database partitions'
     task create_dynamic_partitions: :environment do
-      Gitlab::Database::Partitioning::PartitionManager.new.sync_partitions
+      Gitlab::Database::Partitioning.sync_partitions
     end
 
     # This is targeted towards deploys and upgrades of GitLab.

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) { described_class.new(models).sync_partitions }
+
+  let(:models) { [double, double] }
+
+  let(:db_name1) { 'db1' }
+  let(:db_name2) { 'db2' }
+
+  let(:config1) { 'config1' }
+  let(:config2) { 'config2' }
+  let(:configurations) { double }
+
+  let(:manager_class) { Gitlab::Database::Partitioning::PartitionManager }
+  let(:manager1) { double('manager 1') }
+  let(:manager2) { double('manager 2') }
+
+  let(:original_config) { ActiveRecord::Base.connection_db_config }
+
+  before do
+    allow(configurations).to receive(:configs_for).with(env_name: Rails.env, name: db_name1).and_return(config1)
+    allow(configurations).to receive(:configs_for).with(env_name: Rails.env, name: db_name2).and_return(config2)
+
+    allow(Gitlab::Database).to receive(:db_config_names).and_return([db_name1, db_name2])
+
+    allow(ActiveRecord::Base).to receive(:configurations).twice.and_return(configurations)
+  end
+
+  it 'syncs model partitions for each database connection' do
+    expect(ActiveRecord::Base).to receive(:establish_connection).with(config1).ordered
+    expect(manager_class).to receive(:new).with(models).and_return(manager1).ordered
+    expect(manager1).to receive(:sync_partitions).ordered
+
+    expect(ActiveRecord::Base).to receive(:establish_connection).with(config2).ordered
+    expect(manager_class).to receive(:new).with(models).and_return(manager2).ordered
+    expect(manager2).to receive(:sync_partitions).ordered
+
+    expect(ActiveRecord::Base).to receive(:establish_connection).with(original_config).ordered
+
+    sync_partitions
+  end
+
+  context 'if an error is raised' do
+    it 'restores the original connection' do
+      expect(ActiveRecord::Base).to receive(:establish_connection).with(config1).ordered
+      expect(manager_class).to receive(:new).with(models).and_return(manager1).ordered
+      expect(manager1).to receive(:sync_partitions).ordered.and_raise(RuntimeError)
+
+      expect(ActiveRecord::Base).to receive(:establish_connection).with(original_config).ordered
+
+      expect { sync_partitions }.to raise_error(RuntimeError)
+    end
+  end
+
+  context 'if no models are given' do
+    let(:models) { [] }
+
+    it 'does nothing, changing no connections' do
+      expect(ActiveRecord::Base).not_to receive(:establish_connection)
+      expect(manager_class).not_to receive(:new)
+
+      sync_partitions
+    end
+  end
+end

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

@@ -12,19 +12,6 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
     end
   end
 
-  describe '.register' do
-    let(:model) { double(partitioning_strategy: nil) }
-
-    it 'remembers registered models' do
-      expect { described_class.register(model) }.to change { described_class.models }.to include(model)
-    end
-
-    after do
-      # Do not leak the double to other specs
-      described_class.models.delete(model)
-    end
-  end
-
   context 'creating partitions (mocked)' do
     subject(:sync_partitions) { described_class.new(models).sync_partitions }
 

spec/lib/gitlab/database/partitioning_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Partitioning do
+  describe '.sync_partitions' do
+    let(:partition_manager_class) { described_class::MultiDatabasePartitionManager }
+    let(:partition_manager) { double('partition manager') }
+
+    context 'when no partitioned models are given' do
+      it 'calls the partition manager with the default partitions' do
+        expect(partition_manager_class).to receive(:new)
+          .with(described_class.default_partitioned_models)
+          .and_return(partition_manager)
+
+        expect(partition_manager).to receive(:sync_partitions)
+
+        described_class.sync_partitions
+      end
+    end
+
+    context 'when partitioned models are given' do
+      it 'calls the partition manager with the given partitions' do
+        models = ['my special model']
+
+        expect(partition_manager_class).to receive(:new)
+          .with(models)
+          .and_return(partition_manager)
+
+        expect(partition_manager).to receive(:sync_partitions)
+
+        described_class.sync_partitions(models)
+      end
+    end
+  end
+
+  describe '.default_partitioned_models' do
+    subject(:default_partitioned_models) { described_class.default_partitioned_models }
+
+    it 'returns all core and EE models' do
+      core_models = described_class.core_partitioned_models
+      ee_models = described_class.ee_partitioned_models
+
+      expect(default_partitioned_models).to eq(core_models.union(ee_models))
+    end
+  end
+end

spec/models/concerns/partitioned_table_spec.rb

@@ -35,11 +35,5 @@ RSpec.describe PartitionedTable do
 
       expect(my_class.partitioning_strategy.partitioning_key).to eq(key)
     end
-
-    it 'registers itself with the PartitionCreator' do
-      expect(Gitlab::Database::Partitioning::PartitionManager).to receive(:register).with(my_class)
-
-      subject
-    end
   end
 end

spec/support/helpers/migrations_helpers.rb

@@ -30,7 +30,7 @@ module MigrationsHelpers
       end
     end
 
-    klass.tap { Gitlab::Database::Partitioning::PartitionManager.new.sync_partitions }
+    klass.tap { Gitlab::Database::Partitioning::PartitionManager.new([klass]).sync_partitions }
   end
 
   def migrations_paths

spec/workers/database/partition_management_worker_spec.rb

@@ -6,16 +6,14 @@ RSpec.describe Database::PartitionManagementWorker do
   describe '#perform' do
     subject { described_class.new.perform }
 
-    let(:manager) {  instance_double('PartitionManager', sync_partitions: nil) }
     let(:monitoring) { instance_double('PartitionMonitoring', report_metrics: nil) }
 
     before do
-      allow(Gitlab::Database::Partitioning::PartitionManager).to receive(:new).and_return(manager)
       allow(Gitlab::Database::Partitioning::PartitionMonitoring).to receive(:new).and_return(monitoring)
     end
 
-    it 'delegates to PartitionManager' do
-      expect(manager).to receive(:sync_partitions)
+    it 'delegates to Partitioning' do
+      expect(Gitlab::Database::Partitioning).to receive(:sync_partitions)
 
       subject
     end