1 files changed, 49 insertions, 1 deletions

diff --git a/lib/gitlab/database/load_balancing/configuration.rb b/lib/gitlab/database/load_balancing/configuration.rb
index 6156515bd73..da313361073 100644
--- a/lib/gitlab/database/load_balancing/configuration.rb
+++ b/lib/gitlab/database/load_balancing/configuration.rb
@@ -7,7 +7,7 @@ module Gitlab
       class Configuration
         attr_accessor :hosts, :max_replication_difference,
                       :max_replication_lag_time, :replica_check_interval,
-                      :service_discovery, :model
+                      :service_discovery
 
         # Creates a configuration object for the given ActiveRecord model.
         def self.for_model(model)
@@ -41,6 +41,8 @@ module Gitlab
             end
           end
 
+          config.reuse_primary_connection!
+
           config
         end
 
@@ -59,6 +61,28 @@ module Gitlab
             disconnect_timeout: 120,
             use_tcp: false
           }
+
+          # Temporary model for GITLAB_LOAD_BALANCING_REUSE_PRIMARY_
+          # To be removed with FF
+          @primary_model = nil
+        end
+
+        def db_config_name
+          @model.connection_db_config.name.to_sym
+        end
+
+        # With connection re-use the primary connection can be overwritten
+        # to be used from different model
+        def primary_connection_specification_name
+          (@primary_model || @model).connection_specification_name
+        end
+
+        def primary_db_config
+          (@primary_model || @model).connection_db_config
+        end
+
+        def replica_db_config
+          @model.connection_db_config
         end
 
         def pool_size
@@ -86,6 +110,30 @@ module Gitlab
         def service_discovery_enabled?
           service_discovery[:record].present?
         end
+
+        # TODO: This is temporary code to allow re-use of primary connection
+        # if the two connections are pointing to the same host. This is needed
+        # to properly support transaction visibility.
+        #
+        # This behavior is required to support [Phase 3](https://gitlab.com/groups/gitlab-org/-/epics/6160#progress).
+        # This method is meant to be removed as soon as it is finished.
+        #
+        # The remapping is done as-is:
+        #   export GITLAB_LOAD_BALANCING_REUSE_PRIMARY_<name-of-connection>=<new-name-of-connection>
+        #
+        # Ex.:
+        #   export GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci=main
+        #
+        def reuse_primary_connection!
+          new_connection = ENV["GITLAB_LOAD_BALANCING_REUSE_PRIMARY_#{db_config_name}"]
+          return unless new_connection.present?
+
+          @primary_model = Gitlab::Database.database_base_models[new_connection.to_sym]
+
+          unless @primary_model
+            raise "Invalid value for 'GITLAB_LOAD_BALANCING_REUSE_PRIMARY_#{db_config_name}=#{new_connection}'"
+          end
+        end
       end
     end
   end