diff options
Diffstat (limited to 'lib/gitlab/database/each_database.rb')
-rw-r--r-- | lib/gitlab/database/each_database.rb | 32 |
1 files changed, 27 insertions, 5 deletions
diff --git a/lib/gitlab/database/each_database.rb b/lib/gitlab/database/each_database.rb index 7c9e65e6691..c3eea0515d4 100644 --- a/lib/gitlab/database/each_database.rb +++ b/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) |