1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
|
# frozen_string_literal: true
# rubocop:disable GitlabSecurity/PublicSend
module Gitlab
module Database
module LoadBalancing
# Redirecting of ActiveRecord connections.
#
# The ConnectionProxy class redirects ActiveRecord connection requests to
# the right load balancer pool, depending on the type of query.
class ConnectionProxy
WriteInsideReadOnlyTransactionError = Class.new(StandardError)
READ_ONLY_TRANSACTION_KEY = :load_balacing_read_only_transaction
attr_reader :load_balancer
# These methods perform writes after which we need to stick to the
# primary.
STICKY_WRITES = %i(
delete
delete_all
insert
update
update_all
).freeze
NON_STICKY_READS = %i(
sanitize_limit
select
select_one
select_rows
quote_column_name
).freeze
# hosts - The hosts to use for load balancing.
def initialize(load_balancer)
@load_balancer = load_balancer
end
def select_all(arel, name = nil, binds = [], preparable: nil)
if arel.respond_to?(:locked) && arel.locked
# SELECT ... FOR UPDATE queries should be sent to the primary.
current_session.write!
write_using_load_balancer(:select_all, arel, name, binds)
else
read_using_load_balancer(:select_all, arel, name, binds)
end
end
NON_STICKY_READS.each do |name|
define_method(name) do |*args, **kwargs, &block|
read_using_load_balancer(name, *args, **kwargs, &block)
end
end
STICKY_WRITES.each do |name|
define_method(name) do |*args, **kwargs, &block|
current_session.write!
write_using_load_balancer(name, *args, **kwargs, &block)
end
end
def transaction(*args, **kwargs, &block)
if current_session.fallback_to_replicas_for_ambiguous_queries?
track_read_only_transaction!
read_using_load_balancer(:transaction, *args, **kwargs, &block)
else
current_session.write!
write_using_load_balancer(:transaction, *args, **kwargs, &block)
end
ensure
untrack_read_only_transaction!
end
def respond_to_missing?(name, include_private = false)
@load_balancer.read_write do |connection|
connection.respond_to?(name, include_private)
end
end
# Delegates all unknown messages to a read-write connection.
def method_missing(...)
if current_session.fallback_to_replicas_for_ambiguous_queries?
read_using_load_balancer(...)
else
write_using_load_balancer(...)
end
end
# Performs a read using the load balancer.
#
# name - The name of the method to call on a connection object.
def read_using_load_balancer(...)
if current_session.use_primary? &&
!current_session.use_replicas_for_read_queries?
@load_balancer.read_write do |connection|
connection.send(...)
end
else
@load_balancer.read do |connection|
connection.send(...)
end
end
end
# Performs a write using the load balancer.
#
# name - The name of the method to call on a connection object.
# sticky - If set to true the session will stick to the master after
# the write.
def write_using_load_balancer(...)
if read_only_transaction?
raise WriteInsideReadOnlyTransactionError, 'A write query is performed inside a read-only transaction'
end
@load_balancer.read_write do |connection|
connection.send(...)
end
end
private
def current_session
::Gitlab::Database::LoadBalancing::Session.current
end
def track_read_only_transaction!
Thread.current[READ_ONLY_TRANSACTION_KEY] = true
end
def untrack_read_only_transaction!
Thread.current[READ_ONLY_TRANSACTION_KEY] = nil
end
def read_only_transaction?
Thread.current[READ_ONLY_TRANSACTION_KEY] == true
end
end
end
end
end
|
