Commit fc2c4e0d authored by Yorick Peterse's avatar Yorick Peterse

Merge branch 'move-get-write-location' into 'master'

Move `get_write_location` into `Database::LoadBalancing`

See merge request gitlab-org/gitlab!72626
parents 3d309089 8cf84c91
......@@ -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
......
......@@ -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
......
......@@ -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
......@@ -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
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment