Move `get_write_location` into `Database::LoadBalancing`

Since the `get_write_location` is used only in `LB` context
this moves it to be only available there.

lib/gitlab/database/connection.rb

@@ -200,33 +200,6 @@ module Gitlab
         row['result'] if row
       end
 
-      # @param [ActiveRecord::Connection] ar_connection
-      # @return [String]
-      def get_write_location(ar_connection)
-        use_new_load_balancer_query = Gitlab::Utils
-          .to_boolean(ENV['USE_NEW_LOAD_BALANCER_QUERY'], default: true)
-
-        sql =
-          if use_new_load_balancer_query
-            <<~NEWSQL
-              SELECT CASE
-                  WHEN pg_is_in_recovery() = true AND EXISTS (SELECT 1 FROM pg_stat_get_wal_senders())
-                    THEN pg_last_wal_replay_lsn()::text
-                  WHEN pg_is_in_recovery() = false
-                    THEN pg_current_wal_insert_lsn()::text
-                    ELSE NULL
-                  END AS location;
-            NEWSQL
-          else
-            <<~SQL
-              SELECT pg_current_wal_insert_lsn()::text AS location
-            SQL
-          end
-
-        row = ar_connection.select_all(sql).first
-        row['location'] if row
-      end
-
       # inside_transaction? will return true if the caller is running within a
       # transaction. Handles special cases when running inside a test
       # environment, where tests may be wrapped in transactions

lib/gitlab/database/load_balancing/load_balancer.rb

@@ -133,7 +133,7 @@ module Gitlab
         # Returns the transaction write location of the primary.
         def primary_write_location
           location = read_write do |connection|
-            ::Gitlab::Database.main.get_write_location(connection)
+            get_write_location(connection)
           end
 
           return location if location
@@ -268,6 +268,33 @@ module Gitlab
           base = SafeRequestStore[:gitlab_load_balancer] ||= {}
           base[self] ||= {}
         end
+
+        # @param [ActiveRecord::Connection] ar_connection
+        # @return [String]
+        def get_write_location(ar_connection)
+          use_new_load_balancer_query = Gitlab::Utils
+            .to_boolean(ENV['USE_NEW_LOAD_BALANCER_QUERY'], default: true)
+
+          sql =
+            if use_new_load_balancer_query
+              <<~NEWSQL
+                SELECT CASE
+                    WHEN pg_is_in_recovery() = true AND EXISTS (SELECT 1 FROM pg_stat_get_wal_senders())
+                      THEN pg_last_wal_replay_lsn()::text
+                    WHEN pg_is_in_recovery() = false
+                      THEN pg_current_wal_insert_lsn()::text
+                      ELSE NULL
+                    END AS location;
+              NEWSQL
+            else
+              <<~SQL
+                SELECT pg_current_wal_insert_lsn()::text AS location
+              SQL
+            end
+
+          row = ar_connection.select_all(sql).first
+          row['location'] if row
+        end
       end
     end
   end

spec/lib/gitlab/database/connection_spec.rb

@@ -428,15 +428,4 @@ RSpec.describe Gitlab::Database::Connection do
       expect(connection.system_id).to be_an_instance_of(Integer)
     end
   end
-
-  describe '#get_write_location' do
-    it 'returns a string' do
-      expect(connection.get_write_location(connection.scope.connection))
-        .to be_a(String)
-    end
-
-    it 'returns nil if there are no results' do
-      expect(connection.get_write_location(double(select_all: []))).to be_nil
-    end
-  end
 end

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

@@ -441,4 +441,15 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
       lb.disconnect!(timeout: 30)
     end
   end
+
+  describe '#get_write_location' do
+    it 'returns a string' do
+      expect(lb.send(:get_write_location, lb.pool.connection))
+        .to be_a(String)
+    end
+
+    it 'returns nil if there are no results' do
+      expect(lb.send(:get_write_location, double(select_all: []))).to be_nil
+    end
+  end
 end