1 files changed, 61 insertions, 0 deletions
diff --git a/lib/gitlab/database/load_balancing/setup.rb b/lib/gitlab/database/load_balancing/setup.rb
new file mode 100644
index 00000000000..3cce839a960
--- /dev/null
+++ b/lib/gitlab/database/load_balancing/setup.rb
@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module LoadBalancing
+      # Class for setting up load balancing of a specific model.
+      class Setup
+        attr_reader :configuration
+
+        def initialize(model, start_service_discovery: false)
+          @model = model
+          @configuration = Configuration.for_model(model)
+          @start_service_discovery = start_service_discovery
+        end
+
+        def setup
+          disable_prepared_statements
+          setup_load_balancer
+          setup_service_discovery
+        end
+
+        def disable_prepared_statements
+          db_config_object = @model.connection_db_config
+          config =
+            db_config_object.configuration_hash.merge(prepared_statements: false)
+
+          hash_config = ActiveRecord::DatabaseConfigurations::HashConfig.new(
+            db_config_object.env_name,
+            db_config_object.name,
+            config
+          )
+
+          @model.establish_connection(hash_config)
+        end
+
+        def setup_load_balancer
+          lb = LoadBalancer.new(configuration)
+
+          # We just use a simple `class_attribute` here so we don't need to
+          # inject any modules and/or expose unnecessary methods.
+          @model.class_attribute(:connection)
+          @model.class_attribute(:sticking)
+
+          @model.connection = ConnectionProxy.new(lb)
+          @model.sticking = Sticking.new(lb)
+        end
+
+        def setup_service_discovery
+          return unless configuration.service_discovery_enabled?
+
+          lb = @model.connection.load_balancer
+          sv = ServiceDiscovery.new(lb, **configuration.service_discovery)
+
+          sv.perform_service_discovery
+
+          sv.start if @start_service_discovery
+        end
+      end
+    end
+  end
+end