Add latest changes from gitlab-org/gitlab@master

1 files changed, 17 insertions, 2 deletions

diff --git a/lib/gitlab/database/each_database.rb b/lib/gitlab/database/each_database.rb
index c3eea0515d4..61d20cfe071 100644
--- a/lib/gitlab/database/each_database.rb
+++ b/lib/gitlab/database/each_database.rb
@@ -4,8 +4,11 @@ module Gitlab
   module Database
     module EachDatabase
       class << self
-        def each_database_connection
-          Gitlab::Database.database_base_models.each_pair do |connection_name, model|
+        def each_database_connection(only: nil)
+          selected_names = Array.wrap(only)
+          base_models = select_base_models(selected_names)
+
+          base_models.each_pair do |connection_name, model|
             connection = model.connection
 
             with_shared_connection(connection, connection_name) do
@@ -28,6 +31,18 @@ module Gitlab
 
         private
 
+        def select_base_models(names)
+          base_models = Gitlab::Database.database_base_models
+
+          return base_models if names.empty?
+
+          names.each_with_object(HashWithIndifferentAccess.new) do |name, hash|
+            raise ArgumentError, "#{name} is not a valid database name" unless base_models.key?(name)
+
+            hash[name] = base_models[name]
+          end
+        end
+
         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