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