Replace hardcoded database config name :main

lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb

@@ -32,12 +32,10 @@ module Gitlab
         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.
+          #      connection.load_balancer.primary_write_location
           #   end
           #
-          # TODO: Replace hardcoded database config name :main when we merge unification strategy
-          # https://gitlab.com/gitlab-org/gitlab/-/issues/336566
-          job['wal_locations'] = { main: wal_location } if wal_location
+          job['wal_locations'] = { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => wal_location } if wal_location
         end
 
         def wal_location

lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb

@@ -66,7 +66,7 @@ module Gitlab
         def legacy_wal_location(job)
           wal_location = job['database_write_location'] || job['database_replica_location']
 
-          { main: wal_location } if wal_location
+          { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => wal_location } if wal_location
         end
 
         def load_balancing_available?(worker_class)

spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb

@@ -85,7 +85,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
         end
 
         it 'passes database_replica_location' do
-          expected_location = { main: location }
+          expected_location = { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => location }
 
           expect(load_balancer).to receive_message_chain(:host, "database_replica_location").and_return(location)
 
@@ -103,7 +103,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
         end
 
         it 'passes primary write location', :aggregate_failures do
-          expected_location = { main: location }
+          expected_location = { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => location }
 
           expect(load_balancer).to receive(:primary_write_location).and_return(location)
 
@@ -117,6 +117,16 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
     end
 
     shared_examples_for 'database location was already provided' do
+      let(:old_location) { '0/D525E3A8' }
+      let(:new_location) { 'AB/12345' }
+      let(:wal_locations) { { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => old_location } }
+      let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'wal_locations' => wal_locations } }
+
+      before do
+        allow(load_balancer).to receive(:primary_write_location).and_return(new_location)
+        allow(load_balancer).to receive(:database_replica_location).and_return(new_location)
+      end
+
       shared_examples_for 'does not set database location again' do |use_primary|
         before do
           allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary?).and_return(use_primary)
@@ -125,19 +135,10 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
         it 'does not set database locations again' do
           run_middleware
 
-          expect(job['wal_locations']).to eq({ main: old_location })
+          expect(job['wal_locations']).to eq(wal_locations)
         end
       end
 
-      let(:old_location) { '0/D525E3A8' }
-      let(:new_location) { 'AB/12345' }
-      let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'wal_locations' => { main: old_location } } }
-
-      before do
-        allow(load_balancer).to receive(:primary_write_location).and_return(new_location)
-        allow(load_balancer).to receive(:database_replica_location).and_return(new_location)
-      end
-
       context "when write was performed" do
         include_examples 'does not set database location again', true
       end

spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb

@@ -63,10 +63,11 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
     end
 
     shared_examples_for 'replica is up to date' do |expected_strategy|
-      let(:wal_locations) { { main: '0/D525E3A8' } }
+      let(:location) {'0/D525E3A8' }
+      let(:wal_locations) { { Gitlab::Database::MAIN_DATABASE_NAME.to_sym => location } }
 
       it 'does not stick to the primary', :aggregate_failures do
-        expect(load_balancer).to receive(:select_up_to_date_host).with(wal_locations[:main]).and_return(true)
+        expect(load_balancer).to receive(:select_up_to_date_host).with(location).and_return(true)
 
         run_middleware do
           expect(Gitlab::Database::LoadBalancing::Session.current.use_primary?).not_to be_truthy
@@ -91,7 +92,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqServerMiddleware do
         let(:job) { { 'job_id' => 'a180b47c-3fd6-41b8-81e9-34da61c3400e', 'wal_locations' => wal_locations } }
 
         before do
-          allow(load_balancer).to receive(:select_up_to_date_host).with(wal_locations[:main]).and_return(true)
+          allow(load_balancer).to receive(:select_up_to_date_host).with(location).and_return(true)
         end
 
         it_behaves_like 'replica is up to date', 'replica'