New approach passing connection into shared models

Rework multi-database partitioning code to not use establish_connection,
but instead use the connection from the partitioned model, which should
already be setup for the correct database.

app/models/postgresql/detached_partition.rb

 # frozen_string_literal: true
 
 module Postgresql
-  class DetachedPartition < ApplicationRecord
+  class DetachedPartition < ::Gitlab::Database::SharedModel
     scope :ready_to_drop, -> { where('drop_after < ?', Time.current) }
   end
 end

config/initializers/postgres_partitioning.rb

 # frozen_string_literal: true
 
+Gitlab::Database::Partitioning.register_models([
+  AuditEvent,
+  WebHookLog,
+  LooseForeignKeys::DeletedRecord
+])
+
+if Gitlab.ee?
+  Gitlab::Database::Partitioning.register_models([
+    IncidentManagement::PendingEscalations::Alert,
+    IncidentManagement::PendingEscalations::Issue
+  ])
+end
+
 begin
   Gitlab::Database::Partitioning.sync_partitions unless ENV['DISABLE_POSTGRES_PARTITION_CREATION_ON_STARTUP']
 rescue ActiveRecord::ActiveRecordError, PG::Error

lib/gitlab/database/partitioning.rb

@@ -3,28 +3,16 @@
 module Gitlab
   module Database
     module Partitioning
-      def self.sync_partitions(partitioned_models = default_partitioned_models)
-        MultiDatabasePartitionManager.new(partitioned_models).sync_partitions
+      def self.register_models(models)
+        registered_models.merge(models)
       end
 
-      def self.default_partitioned_models
-        @default_partitioned_models ||= core_partitioned_models.union(ee_partitioned_models)
+      def self.registered_models
+        @registered_models ||= Set.new
       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
+      def self.sync_partitions(models_to_sync = registered_models)
+        MultiDatabasePartitionManager.new(models_to_sync).sync_partitions
       end
     end
   end

lib/gitlab/database/partitioning/multi_database_partition_manager.rb

@@ -9,35 +9,27 @@ module Gitlab
         end
 
         def sync_partitions
-          return if models.empty?
+          Gitlab::AppLogger.info(message: "Syncing dynamic postgres partitions")
 
-          each_database_connection do
-            PartitionManager.new(models).sync_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
 
-        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
+          Gitlab::AppLogger.info(message: "Finished sync of dynamic postgres partitions")
         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
+        private
 
-            ActiveRecord::Base.establish_connection(config_for_db_name)
+        attr_reader :models
 
-            yield
-          end
+        def connection_name
+          Gitlab::Database::SharedModel.connection.pool.db_config.name
         end
       end
     end

lib/gitlab/database/partitioning/partition_manager.rb

@@ -10,24 +10,23 @@ module Gitlab
         MANAGEMENT_LEASE_KEY = 'database_partition_management_%s'
         RETAIN_DETACHED_PARTITIONS_FOR = 1.week
 
-        def initialize(models)
-          @models = models
+        def initialize(model)
+          @model = model
         end
 
         def sync_partitions
-          Gitlab::AppLogger.info("Checking state of dynamic postgres partitions")
+          Gitlab::AppLogger.info(message: "Checking state of dynamic postgres partitions", table_name: model.table_name)
 
-          models.each do |model|
           # Double-checking before getting the lease:
           # The prevailing situation is no missing partitions and no extra partitions
-            next if missing_partitions(model).empty? && extra_partitions(model).empty?
+          return if missing_partitions.empty? && extra_partitions.empty?
 
           only_with_exclusive_lease(model, lease_key: MANAGEMENT_LEASE_KEY) do
-              partitions_to_create = missing_partitions(model)
+            partitions_to_create = missing_partitions
             create(partitions_to_create) unless partitions_to_create.empty?
 
             if Feature.enabled?(:partition_pruning, default_enabled: :yaml)
-                partitions_to_detach = extra_partitions(model)
+              partitions_to_detach = extra_partitions
               detach(partitions_to_detach) unless partitions_to_detach.empty?
             end
           end
@@ -37,19 +36,19 @@ module Gitlab
                                   exception_class: e.class,
                                   exception_message: e.message)
         end
-        end
 
         private
 
-        attr_reader :models
+        attr_reader :model
+        delegate :connection, to: :model
 
-        def missing_partitions(model)
+        def missing_partitions
           return [] unless connection.table_exists?(model.table_name)
 
           model.partitioning_strategy.missing_partitions
         end
 
-        def extra_partitions(model)
+        def extra_partitions
           return [] unless connection.table_exists?(model.table_name)
 
           model.partitioning_strategy.extra_partitions
@@ -111,13 +110,10 @@ module Gitlab
         def with_lock_retries(&block)
           Gitlab::Database::WithLockRetries.new(
             klass: self.class,
-            logger: Gitlab::AppLogger
+            logger: Gitlab::AppLogger,
+            connection: connection
           ).run(&block)
         end
-
-        def connection
-          ActiveRecord::Base.connection
-        end
       end
     end
   end

lib/gitlab/database/partitioning/partition_monitoring.rb

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

lib/gitlab/database/postgres_foreign_key.rb

@@ -2,7 +2,7 @@
 
 module Gitlab
   module Database
-    class PostgresForeignKey < ApplicationRecord
+    class PostgresForeignKey < SharedModel
       self.primary_key = :oid
 
       scope :by_referenced_table_identifier, ->(identifier) do

lib/gitlab/database/postgres_partition.rb

@@ -2,7 +2,7 @@
 
 module Gitlab
   module Database
-    class PostgresPartition < ActiveRecord::Base
+    class PostgresPartition < SharedModel
       self.primary_key = :identifier
 
       belongs_to :postgres_partitioned_table, foreign_key: 'parent_identifier', primary_key: 'identifier'

lib/gitlab/database/postgres_partitioned_table.rb

@@ -2,7 +2,7 @@
 
 module Gitlab
   module Database
-    class PostgresPartitionedTable < ActiveRecord::Base
+    class PostgresPartitionedTable < SharedModel
       DYNAMIC_PARTITION_STRATEGIES = %w[range list].freeze
 
       self.primary_key = :identifier

lib/gitlab/database/shared_model.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    class SharedModel < ActiveRecord::Base
+      self.abstract_class = true
+
+      class << self
+        def using_connection(connection)
+          raise 'cannot nest connection overrides for shared models' unless overriding_connection.nil?
+
+          self.overriding_connection = connection
+
+          yield
+        ensure
+          self.overriding_connection = nil
+        end
+
+        def connection
+          if connection = self.overriding_connection
+            connection
+          else
+            super
+          end
+        end
+
+        private
+
+        def overriding_connection
+          Thread.current[:overriding_connection]
+        end
+
+        def overriding_connection=(connection)
+          Thread.current[:overriding_connection] = connection
+        end
+      end
+    end
+  end
+end

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

@@ -3,66 +3,34 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::Partitioning::MultiDatabasePartitionManager, '#sync_partitions' do
-  subject(:sync_partitions) { described_class.new(models).sync_partitions }
+  subject(:sync_partitions) { manager.sync_partitions }
 
-  let(:models) { [double, double] }
+  let(:manager) { described_class.new(models) }
+  let(:models) { [model1, model2] }
 
-  let(:db_name1) { 'db1' }
-  let(:db_name2) { 'db2' }
+  let(:model1) { double('model1', connection: connection1, table_name: 'table1') }
+  let(:model2) { double('model2', connection: connection1, table_name: 'table2') }
 
-  let(:config1) { 'config1' }
-  let(:config2) { 'config2' }
-  let(:configurations) { double }
+  let(:connection1) { double('connection1') }
+  let(:connection2) { double('connection2') }
 
-  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 }
+  let(:target_manager_class) { Gitlab::Database::Partitioning::PartitionManager }
+  let(:target_manager1) { double('partition manager') }
+  let(:target_manager2) { double('partition manager') }
 
   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)
+    allow(manager).to receive(:connection_name).and_return('name')
   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
+  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(ActiveRecord::Base).to receive(:establish_connection).with(original_config).ordered
+    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
-
-  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

@@ -13,17 +13,17 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
   end
 
   context 'creating partitions (mocked)' do
-    subject(:sync_partitions) { described_class.new(models).sync_partitions }
+    subject(:sync_partitions) { described_class.new(model).sync_partitions }
 
-    let(:models) { [model] }
-    let(:model) { double(partitioning_strategy: partitioning_strategy, table_name: table) }
+    let(:model) { double(partitioning_strategy: partitioning_strategy, table_name: table, connection: connection) }
     let(:partitioning_strategy) { double(missing_partitions: partitions, extra_partitions: []) }
+    let(:connection) { ActiveRecord::Base.connection }
     let(:table) { "some_table" }
 
     before do
-      allow(ActiveRecord::Base.connection).to receive(:table_exists?).and_call_original
-      allow(ActiveRecord::Base.connection).to receive(:table_exists?).with(table).and_return(true)
-      allow(ActiveRecord::Base.connection).to receive(:execute).and_call_original
+      allow(connection).to receive(:table_exists?).and_call_original
+      allow(connection).to receive(:table_exists?).with(table).and_return(true)
+      allow(connection).to receive(:execute).and_call_original
 
       stub_exclusive_lease(described_class::MANAGEMENT_LEASE_KEY % table, timeout: described_class::LEASE_TIMEOUT)
     end
@@ -36,35 +36,23 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
     end
 
     it 'creates the partition' do
-      expect(ActiveRecord::Base.connection).to receive(:execute).with(partitions.first.to_sql)
-      expect(ActiveRecord::Base.connection).to receive(:execute).with(partitions.second.to_sql)
+      expect(connection).to receive(:execute).with(partitions.first.to_sql)
+      expect(connection).to receive(:execute).with(partitions.second.to_sql)
 
       sync_partitions
     end
 
-    context 'error handling with 2 models' do
-      let(:models) do
-        [
-          double(partitioning_strategy: strategy1, table_name: table),
-          double(partitioning_strategy: strategy2, table_name: table)
-        ]
-      end
-
-      let(:strategy1) { double('strategy1', missing_partitions: nil, extra_partitions: []) }
-      let(:strategy2) { double('strategy2', missing_partitions: partitions, extra_partitions: []) }
+    context 'when an error occurs during partition management' do
+      it 'does not raise an error' do
+        expect(partitioning_strategy).to receive(:missing_partitions).and_raise('this should never happen (tm)')
 
-      it 'still creates partitions for the second table' do
-        expect(strategy1).to receive(:missing_partitions).and_raise('this should never happen (tm)')
-        expect(ActiveRecord::Base.connection).to receive(:execute).with(partitions.first.to_sql)
-        expect(ActiveRecord::Base.connection).to receive(:execute).with(partitions.second.to_sql)
-
-        sync_partitions
+        expect { sync_partitions }.not_to raise_error
       end
     end
   end
 
   context 'creating partitions' do
-    subject(:sync_partitions) { described_class.new([my_model]).sync_partitions }
+    subject(:sync_partitions) { described_class.new(my_model).sync_partitions }
 
     let(:connection) { ActiveRecord::Base.connection }
     let(:my_model) do
@@ -93,15 +81,15 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
   context 'detaching partitions (mocked)' do
     subject(:sync_partitions) { manager.sync_partitions }
 
-    let(:manager) { described_class.new(models) }
-    let(:models) { [model] }
-    let(:model) { double(partitioning_strategy: partitioning_strategy, table_name: table)}
+    let(:manager) { described_class.new(model) }
+    let(:model) { double(partitioning_strategy: partitioning_strategy, table_name: table, connection: connection) }
     let(:partitioning_strategy) { double(extra_partitions: extra_partitions, missing_partitions: []) }
+    let(:connection) { ActiveRecord::Base.connection }
     let(:table) { "foo" }
 
     before do
-      allow(ActiveRecord::Base.connection).to receive(:table_exists?).and_call_original
-      allow(ActiveRecord::Base.connection).to receive(:table_exists?).with(table).and_return(true)
+      allow(connection).to receive(:table_exists?).and_call_original
+      allow(connection).to receive(:table_exists?).with(table).and_return(true)
 
       stub_exclusive_lease(described_class::MANAGEMENT_LEASE_KEY % table, timeout: described_class::LEASE_TIMEOUT)
     end
@@ -123,24 +111,6 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
 
         sync_partitions
       end
-
-      context 'error handling' do
-        let(:models) do
-          [
-            double(partitioning_strategy: error_strategy, table_name: table),
-            model
-          ]
-        end
-
-        let(:error_strategy) { double(extra_partitions: nil, missing_partitions: []) }
-
-        it 'still drops partitions for the other model' do
-          expect(error_strategy).to receive(:extra_partitions).and_raise('injected error!')
-          extra_partitions.each { |p| expect(manager).to receive(:detach_one_partition).with(p) }
-
-          sync_partitions
-        end
-      end
     end
 
     context 'with the partition_pruning feature flag disabled' do
@@ -163,7 +133,7 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
       end
     end
 
-    subject { described_class.new([my_model]).sync_partitions }
+    subject { described_class.new(my_model).sync_partitions }
 
     let(:connection) { ActiveRecord::Base.connection }
     let(:my_model) do
@@ -272,11 +242,11 @@ RSpec.describe Gitlab::Database::Partitioning::PartitionManager do
     it 'creates partitions for the future then drops the oldest one after a month' do
       # 1 month for the current month, 1 month for the old month that we're retaining data for, headroom
       expected_num_partitions = (Gitlab::Database::Partitioning::MonthlyStrategy::HEADROOM + 2.months) / 1.month
-      expect { described_class.new([my_model]).sync_partitions }.to change { num_partitions(my_model) }.from(0).to(expected_num_partitions)
+      expect { described_class.new(my_model).sync_partitions }.to change { num_partitions(my_model) }.from(0).to(expected_num_partitions)
 
       travel 1.month
 
-      expect { described_class.new([my_model]).sync_partitions }.to change { has_partition(my_model, 2.months.ago.beginning_of_month) }.from(true).to(false).and(change { num_partitions(my_model) }.by(0))
+      expect { described_class.new(my_model).sync_partitions }.to change { has_partition(my_model, 2.months.ago.beginning_of_month) }.from(true).to(false).and(change { num_partitions(my_model) }.by(0))
     end
   end
 end

spec/lib/gitlab/database/partitioning_spec.rb

@@ -8,9 +8,9 @@ RSpec.describe Gitlab::Database::Partitioning do
     let(:partition_manager) { double('partition manager') }
 
     context 'when no partitioned models are given' do
-      it 'calls the partition manager with the default partitions' do
+      it 'calls the partition manager with the registered models' do
         expect(partition_manager_class).to receive(:new)
-          .with(described_class.default_partitioned_models)
+          .with(described_class.registered_models)
           .and_return(partition_manager)
 
         expect(partition_manager).to receive(:sync_partitions)
@@ -20,7 +20,7 @@ RSpec.describe Gitlab::Database::Partitioning do
     end
 
     context 'when partitioned models are given' do
-      it 'calls the partition manager with the given partitions' do
+      it 'calls the partition manager with the given models' do
         models = ['my special model']
 
         expect(partition_manager_class).to receive(:new)
@@ -33,15 +33,4 @@ RSpec.describe Gitlab::Database::Partitioning do
       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/lib/gitlab/database/shared_model_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::SharedModel do
+  describe 'using an external connection' do
+    let!(:original_connection) { described_class.connection }
+    let(:new_connection) { double('connection') }
+
+    it 'overrides the connection for the duration of the block', :aggregate_failures do
+      expect_original_connection_around do
+        described_class.using_connection(new_connection) do
+          expect(described_class.connection).to be(new_connection)
+        end
+      end
+    end
+
+    it 'does not affect connections in other threads', :aggregate_failures do
+      expect_original_connection_around do
+        described_class.using_connection(new_connection) do
+          expect(described_class.connection).to be(new_connection)
+
+          Thread.new do
+            expect(described_class.connection).not_to be(new_connection)
+          end.join
+        end
+      end
+    end
+
+    context 'when the block raises an error', :aggregate_failures do
+      it 're-raises the error, removing the overridden connection' do
+        expect_original_connection_around do
+          expect do
+            described_class.using_connection(new_connection) do
+              expect(described_class.connection).to be(new_connection)
+
+              raise 'here comes an error!'
+            end
+          end.to raise_error(RuntimeError, 'here comes an error!')
+        end
+      end
+    end
+
+    def expect_original_connection_around
+      # For safety, ensure our original connection is distinct from our double
+      # This should be the case, but in case of something leaking we should verify
+      expect(original_connection).not_to be(new_connection)
+      expect(described_class.connection).to be(original_connection)
+
+      yield
+
+      expect(described_class.connection).to be(original_connection)
+    end
+  end
+end

spec/support/helpers/migrations_helpers.rb

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