Merge branch '325291-remove-hardcoded-main-db-from-lb' into 'master'

Replace hardcoded database config name :main

See merge request gitlab-org/gitlab!69551

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)
 
@@ -116,28 +116,43 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
       end
     end
 
-    shared_examples_for 'database location was already provided' do
-      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)
-        end
+    context 'when worker cannot be constantized' do
+      let(:worker_class) { 'ActionMailer::MailDeliveryJob' }
+      let(:expected_consistency) { :always }
 
-        it 'does not set database locations again' do
-          run_middleware
+      include_examples 'does not pass database locations'
+    end
 
-          expect(job['wal_locations']).to eq({ main: old_location })
-        end
-      end
+    context 'when worker class does not include ApplicationWorker' do
+      let(:worker_class) { ActiveJob::QueueAdapters::SidekiqAdapter::JobWrapper }
+      let(:expected_consistency) { :always }
+
+      include_examples 'does not pass database locations'
+    end
 
+    context 'database wal location was already provided' do
       let(:old_location) { '0/D525E3A8' }
       let(:new_location) { 'AB/12345' }
-      let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", 'wal_locations' => { main: old_location } } }
+      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)
+        end
+
+        it 'does not set database locations again' do
+          run_middleware
+
+          expect(job['wal_locations']).to eq(wal_locations)
+        end
+      end
+
       context "when write was performed" do
         include_examples 'does not set database location again', true
       end
@@ -147,24 +162,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
       end
     end
 
-    context 'when worker cannot be constantized' do
-      let(:worker_class) { 'ActionMailer::MailDeliveryJob' }
-      let(:expected_consistency) { :always }
-
-      include_examples 'does not pass database locations'
-    end
-
-    context 'when worker class does not include ApplicationWorker' do
-      let(:worker_class) { ActiveJob::QueueAdapters::SidekiqAdapter::JobWrapper }
-      let(:expected_consistency) { :always }
-
-      include_examples 'does not pass database locations'
-    end
-
-    context 'database wal location was already provided' do
-      include_examples 'database location was already provided'
-    end
-
     context 'when worker data consistency is :always' do
       include_context 'data consistency worker class', :always, :load_balancing_for_test_data_consistency_worker
 

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'