1 files changed, 15 insertions, 15 deletions
diff --git a/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb b/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
index 518a812b406..62dfe75a851 100644
--- a/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
+++ b/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
@@ -30,26 +30,26 @@ module Gitlab
         end
 
         def set_data_consistency_locations!(job)
-          # Once we add support for multiple databases to our load balancer, we would use something like this:
-          #   job['wal_locations'] = Gitlab::Database::DATABASES.transform_values do |connection|
-          #      connection.load_balancer.primary_write_location
-          #   end
-          #
-          job['wal_locations'] = { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => wal_location } if wal_location
-        end
+          locations = {}
 
-        def wal_location
-          strong_memoize(:wal_location) do
-            if Session.current.use_primary?
-              load_balancer.primary_write_location
-            else
-              load_balancer.host.database_replica_location
+          ::Gitlab::Database::LoadBalancing.each_load_balancer do |lb|
+            if (location = wal_location_for(lb))
+              locations[lb.name] = location
             end
           end
+
+          job['wal_locations'] = locations
         end
 
-        def load_balancer
-          LoadBalancing.proxy.load_balancer
+        def wal_location_for(load_balancer)
+          # When only using the primary there's no need for any WAL queries.
+          return if load_balancer.primary_only?
+
+          if ::Gitlab::Database::LoadBalancing::Session.current.use_primary?
+            load_balancer.primary_write_location
+          else
+            load_balancer.host.database_replica_location
+          end
         end
       end
     end