Commit 6b6f6862 authored by Dylan Griffith's avatar Dylan Griffith

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

Support partitioning of `SharedModels`

See merge request gitlab-org/gitlab!79441
parents fc4c1763 a09a1638
......@@ -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
}
......
......@@ -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)
......
......@@ -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
......
......@@ -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)
......
......@@ -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
......
......@@ -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
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment