Commit 1802e343 authored by Quang-Minh Nguyen's avatar Quang-Minh Nguyen

Address feedback in the reviews

Issue https://gitlab.com/gitlab-org/gitlab/-/issues/322133
parent e693ec99
......@@ -11,11 +11,7 @@ module EE
override :with_fast_read_statement_timeout
def with_fast_read_statement_timeout(timeout_ms = 5000)
::Gitlab::Database::LoadBalancing::Session.current.use_replica_if_possible do
transaction(requires_new: true) do
connection.exec_query("SET LOCAL statement_timeout = #{timeout_ms}")
yield
end
super
end
end
end
......
......@@ -62,15 +62,11 @@ module Gitlab
# - The current session already performed writes
# - It prefers to use primary, aka, use_primary or use_primary! were called
def use_replica_if_possible(&blk)
return yield if use_primary? || performed_write?
begin
used_replica = @use_replica
@use_replica = true
yield
ensure
@use_replica = used_replica && !use_primary? && !performed_write?
end
used_replica = @use_replica
@use_replica = true
yield
ensure
@use_replica = used_replica
end
def use_replica?
......
......@@ -409,14 +409,27 @@ RSpec.describe Gitlab::Database::LoadBalancing do
# of the load balancer.
# - A real model with a table backed behind is defined
# - The load balancing module is set up for this module only, as to prevent
# breaking other tests. The replica configuraiton is cloned from the test
# breaking other tests. The replica configuration is cloned from the test
# configuraiton.
# - In each test, we listen to the SQL queries (via sql.active_record
# instrumentaiton) while triggering real queries from the defined model.
# instrumentation) while triggering real queries from the defined model.
# - We assert the desinations (replica/primary) of the queries in order.
describe 'LoadBalancing integration tests', :delete do
before(:all) do
ActiveRecord::Schema.define do
create_table :load_balancing_test, force: true do |t|
t.string :name, null: true
end
end
end
after(:all) do
ActiveRecord::Schema.define do
drop_table :load_balancing_test, force: true
end
end
shared_context 'LoadBalancing setup' do
let!(:license) { create(:license, plan: ::License::PREMIUM_PLAN) }
let(:hosts) { [ActiveRecord::Base.configurations["development"]['host']] }
let(:model) do
Class.new(ApplicationRecord) do
......@@ -425,11 +438,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
end
before do
ActiveRecord::Schema.define do
create_table :load_balancing_test, force: true do |t|
t.string :name, null: true
end
end
stub_licensed_features(db_load_balancing: true)
# Preloading testing class
model.singleton_class.prepend ::Gitlab::Database::LoadBalancing::ActiveRecordProxy
......@@ -446,9 +455,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
after do
subject.clear_configuration
ActiveRecord::Schema.define do
drop_table :load_balancing_test, force: true
end
end
end
......@@ -531,17 +537,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
true, [:primary, :primary, :primary, :primary]
],
# Read-only transaction
[
-> {
model.transaction do
model.first
model.where(name: 'test1').to_a
end
},
true, [:primary, :primary, :primary, :primary]
],
# use_primary
[
-> {
......
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