1 files changed, 8 insertions, 1 deletions
diff --git a/config/initializers/load_balancing.rb b/config/initializers/load_balancing.rb
index f3c977893ba..2b58ae0f642 100644
--- a/config/initializers/load_balancing.rb
+++ b/config/initializers/load_balancing.rb
@@ -1,12 +1,19 @@
 # frozen_string_literal: true
 
+ActiveRecord::Base.singleton_class.attr_accessor :load_balancing_proxy
+
 if Gitlab::Database::LoadBalancing.enable?
-  Gitlab::Database.disable_prepared_statements
+  Gitlab::Database.main.disable_prepared_statements
 
   Gitlab::Application.configure do |config|
     config.middleware.use(Gitlab::Database::LoadBalancing::RackMiddleware)
   end
 
+  # This hijacks the "connection" method to ensure both
+  # `ActiveRecord::Base.connection` and all models use the same load
+  # balancing proxy.
+  ActiveRecord::Base.singleton_class.prepend(Gitlab::Database::LoadBalancing::ActiveRecordProxy)
+
   Gitlab::Database::LoadBalancing.configure_proxy
 
   # This needs to be executed after fork of clustered processes