Don't update db WAL location if already provided

ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb

@@ -20,6 +20,8 @@ module Gitlab
           return unless worker_class.include?(::ApplicationWorker)
           return unless worker_class.get_data_consistency_feature_flag_enabled?
 
+          return if location_already_provided?(job)
+
           job['worker_data_consistency'] = worker_class.get_data_consistency
 
           return if worker_class.get_data_consistency == :always
@@ -31,6 +33,10 @@ module Gitlab
           end
         end
 
+        def location_already_provided?(job)
+          job['database_replica_location'] || job['database_write_location']
+        end
+
         def load_balancer
           LoadBalancing.proxy.load_balancer
         end

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

@@ -82,6 +82,38 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
       end
     end
 
+    shared_examples_for 'database location was already provided' do |provided_database_location, other_location|
+      shared_examples_for 'does not set database location again' do |write_performed|
+        before do
+          allow(Gitlab::Database::LoadBalancing::Session.current).to receive(:performed_write?).and_return(write_performed)
+        end
+
+        it 'does not set database locations again' do
+          middleware.call(worker_class, job, double(:queue), redis_pool) { 10 }
+
+          expect(job[provided_database_location]).to eq(old_location)
+          expect(job[other_location]).to be_nil
+        end
+      end
+
+      let(:old_location) { '0/D525E3A8' }
+      let(:new_location) { 'AB/12345' }
+      let(:job) { { "job_id" => "a180b47c-3fd6-41b8-81e9-34da61c3400e", provided_database_location => old_location } }
+
+      before do
+        allow(middleware).to receive_message_chain(:load_balancer, :primary_write_location).and_return(new_location)
+        allow(middleware).to receive_message_chain(:load_balancer, :database_replica_location).and_return(new_location)
+      end
+
+      context "when write was performed" do
+        include_examples 'does not set database location again', true
+      end
+
+      context "when write was not performed" do
+        include_examples 'does not set database location again', false
+      end
+    end
+
     let(:queue) { 'default' }
     let(:redis_pool) { Sidekiq.redis_pool }
     let(:worker_class) { 'TestDataConsistencyWorker' }
@@ -104,6 +136,14 @@ RSpec.describe Gitlab::Database::LoadBalancing::SidekiqClientMiddleware do
       include_examples 'does not pass database locations'
     end
 
+    context 'database write location was already provided' do
+      include_examples 'database location was already provided', 'database_write_location', 'database_replica_location'
+    end
+
+    context 'database replica location was already provided' do
+      include_examples 'database location was already provided', 'database_replica_location', 'database_write_location'
+    end
+
     context 'when worker data consistency is :always' do
       include_context 'data consistency worker class', :always, :load_balancing_for_test_data_consistency_worker