Merge branch 'fix-partitioning-for-shared-models' into 'master'

Support partitioning of `SharedModels`

See merge request gitlab-org/gitlab!79441

config/initializers/postgres_partitioning.rb

@@ -15,10 +15,12 @@ if Gitlab.ee?
 else
   Gitlab::Database::Partitioning.register_tables([
     {
+      limit_connection_names: %i[main],
       table_name: 'incident_management_pending_alert_escalations',
       partitioned_column: :process_at, strategy: :monthly
     },
     {
+      limit_connection_names: %i[main],
       table_name: 'incident_management_pending_issue_escalations',
       partitioned_column: :process_at, strategy: :monthly
     }
@@ -31,6 +33,7 @@ unless Gitlab.jh?
     # This should be synchronized with the following model:
     # https://jihulab.com/gitlab-cn/gitlab/-/blob/main-jh/jh/app/models/phone/verification_code.rb
     {
+      limit_connection_names: %i[main],
       table_name: 'verification_codes',
       partitioned_column: :created_at, strategy: :monthly
     }

lib/gitlab/database/each_database.rb

@@ -14,18 +14,40 @@ module Gitlab
           end
         end
 
-        def each_model_connection(models)
+        def each_model_connection(models, &blk)
           models.each do |model|
-            connection_name = model.connection.pool.db_config.name
-
-            with_shared_connection(model.connection, connection_name) do
-              yield model, connection_name
+            # If model is shared, iterate all available base connections
+            # Example: `LooseForeignKeys::DeletedRecord`
+            if model < ::Gitlab::Database::SharedModel
+              with_shared_model_connections(model, &blk)
+            else
+              with_model_connection(model, &blk)
             end
           end
         end
 
         private
 
+        def with_shared_model_connections(shared_model, &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
+
+            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
+
+          with_shared_connection(model.connection, connection_name) do
+            yield model, connection_name
+          end
+        end
+
         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)

lib/gitlab/database/partitioning.rb

@@ -3,19 +3,8 @@
 module Gitlab
   module Database
     module Partitioning
-      class TableWithoutModel
-        include PartitionedTable::ClassMethods
-
-        attr_reader :table_name
-
-        def initialize(table_name:, partitioned_column:, strategy:)
-          @table_name = table_name
-          partitioned_by(partitioned_column, strategy: strategy)
-        end
-
-        def connection
-          Gitlab::Database::SharedModel.connection
-        end
+      class TableWithoutModel < Gitlab::Database::SharedModel
+        include PartitionedTable
       end
 
       class << self
@@ -77,7 +66,15 @@ module Gitlab
 
         def registered_for_sync
           registered_models + registered_tables.map do |table|
-            TableWithoutModel.new(**table)
+            table_without_model(**table)
+          end
+        end
+
+        def table_without_model(table_name:, partitioned_column:, strategy:, limit_connection_names: nil)
+          Class.new(TableWithoutModel).tap do |klass|
+            klass.table_name = table_name
+            klass.partitioned_by(partitioned_column, strategy: strategy)
+            klass.limit_connection_names = limit_connection_names
           end
         end
       end

lib/gitlab/database/partitioning/partition_manager.rb

@@ -12,10 +12,15 @@ module Gitlab
 
         def initialize(model)
           @model = model
+          @connection_name = model.connection.pool.db_config.name
         end
 
         def sync_partitions
-          Gitlab::AppLogger.info(message: "Checking state of dynamic postgres partitions", table_name: model.table_name)
+          Gitlab::AppLogger.info(
+            message: "Checking state of dynamic postgres partitions",
+            table_name: model.table_name,
+            connection_name: @connection_name
+          )
 
           # Double-checking before getting the lease:
           # The prevailing situation is no missing partitions and no extra partitions
@@ -29,10 +34,13 @@ module Gitlab
             detach(partitions_to_detach) unless partitions_to_detach.empty?
           end
         rescue StandardError => e
-          Gitlab::AppLogger.error(message: "Failed to create / detach partition(s)",
-                                  table_name: model.table_name,
-                                  exception_class: e.class,
-                                  exception_message: e.message)
+          Gitlab::AppLogger.error(
+            message: "Failed to create / detach partition(s)",
+            table_name: model.table_name,
+            exception_class: e.class,
+            exception_message: e.message,
+            connection_name: @connection_name
+          )
         end
 
         private
@@ -98,9 +106,12 @@ module Gitlab
           Postgresql::DetachedPartition.create!(table_name: partition.partition_name,
                                                 drop_after: RETAIN_DETACHED_PARTITIONS_FOR.from_now)
 
-          Gitlab::AppLogger.info(message: "Detached Partition",
-                                 partition_name: partition.partition_name,
-                                 table_name: partition.table)
+          Gitlab::AppLogger.info(
+            message: "Detached Partition",
+            partition_name: partition.partition_name,
+            table_name: partition.table,
+            connection_name: @connection_name
+          )
         end
 
         def assert_partition_detachable!(partition)

lib/gitlab/database/shared_model.rb

@@ -6,6 +6,10 @@ module Gitlab
     class SharedModel < ActiveRecord::Base
       self.abstract_class = true
 
+      # if shared model is used, this allows to limit connections
+      # on which this model is being shared
+      class_attribute :limit_connection_names, default: nil
+
       class << self
         def using_connection(connection)
           previous_connection = self.overriding_connection

spec/lib/gitlab/database/each_database_spec.rb

@@ -4,45 +4,97 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Database::EachDatabase do
   describe '.each_database_connection' do
-    let(:expected_connections) do
-      Gitlab::Database.database_base_models.map { |name, model| [model.connection, name] }
+    before do
+      allow(Gitlab::Database).to receive(:database_base_models)
+        .and_return({ main: ActiveRecord::Base, ci: Ci::ApplicationRecord }.with_indifferent_access)
     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
+    it 'yields each connection after connecting SharedModel', :add_ci_connection do
+      expect(Gitlab::Database::SharedModel).to receive(:using_connection)
+        .with(ActiveRecord::Base.connection).ordered.and_yield
 
-      yielded_connections = []
+      expect(Gitlab::Database::SharedModel).to receive(:using_connection)
+        .with(Ci::ApplicationRecord.connection).ordered.and_yield
 
-      described_class.each_database_connection do |connection, name|
-        yielded_connections << [connection, name]
-      end
-
-      expect(yielded_connections).to match_array(expected_connections)
+      expect { |b| described_class.each_database_connection(&b) }
+        .to yield_successive_args(
+          [ActiveRecord::Base.connection, 'main'],
+          [Ci::ApplicationRecord.connection, 'ci']
+        )
     end
   end
 
   describe '.each_model_connection' do
-    let(:model1) { double(connection: double, table_name: 'table1') }
-    let(:model2) { double(connection: double, table_name: 'table2') }
+    context 'when the model inherits from SharedModel', :add_ci_connection do
+      let(:model1) { Class.new(Gitlab::Database::SharedModel) }
+      let(:model2) { Class.new(Gitlab::Database::SharedModel) }
 
-    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')
+      before do
+        allow(Gitlab::Database).to receive(:database_base_models)
+          .and_return({ main: ActiveRecord::Base, ci: Ci::ApplicationRecord }.with_indifferent_access)
+      end
+
+      it 'yields each model with SharedModel connected to each database connection' do
+        expect_yielded_models([model1, model2], [
+          { model: model1, connection: ActiveRecord::Base.connection, name: 'main' },
+          { model: model1, connection: Ci::ApplicationRecord.connection, name: 'ci' },
+          { model: model2, connection: ActiveRecord::Base.connection, name: 'main' },
+          { model: model2, connection: Ci::ApplicationRecord.connection, name: 'ci' }
+        ])
+      end
+
+      context 'when the model limits connection names' do
+        before do
+          model1.limit_connection_names = %i[main]
+          model2.limit_connection_names = %i[ci]
+        end
+
+        it 'only yields the model with SharedModel connected to the limited connections' do
+          expect_yielded_models([model1, model2], [
+            { model: model1, connection: ActiveRecord::Base.connection, name: 'main' },
+            { model: model2, connection: Ci::ApplicationRecord.connection, name: 'ci' }
+          ])
+        end
+      end
+    end
+
+    context 'when the model does not inherit from SharedModel' do
+      let(:main_model) { Class.new(ActiveRecord::Base) }
+      let(:ci_model) { Class.new(Ci::ApplicationRecord) }
+
+      let(:main_connection) { double(:connection) }
+      let(:ci_connection) { double(:connection) }
+
+      before 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')
+      end
+
+      it 'yields each model after connecting SharedModel' do
+        expect_yielded_models([main_model, ci_model], [
+          { model: main_model, connection: main_connection, name: 'main' },
+          { model: ci_model, connection: ci_connection, name: 'ci' }
+        ])
+      end
     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
+    def expect_yielded_models(models_to_iterate, expected_values)
+      times_yielded = 0
+
+      described_class.each_model_connection(models_to_iterate) do |model, name|
+        expected = expected_values[times_yielded]
 
-      yielded_models = []
+        expect(model).to be(expected[:model])
+        expect(model.connection).to be(expected[:connection])
+        expect(name).to eq(expected[:name])
 
-      described_class.each_model_connection([model1, model2]) do |model, name|
-        yielded_models << [model, name]
+        times_yielded += 1
       end
 
-      expect(yielded_models).to match_array([[model1, 'name1'], [model2, 'name2']])
+      expect(times_yielded).to eq(expected_values.size)
     end
   end
 end