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
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
|
# frozen_string_literal: true
module Gitlab
module Database
# Configuration settings and methods for interacting with a PostgreSQL
# database, with support for multiple databases.
class Connection
attr_reader :scope
# Initializes a new `Database`.
#
# The `scope` argument must be an object (such as `ActiveRecord::Base`)
# that supports retrieving connections and connection pools.
def initialize(scope = ActiveRecord::Base)
@config = nil
@scope = scope
@version = nil
@open_transactions_baseline = 0
end
def config
# The result of this method must not be cached, as other methods may use
# it after making configuration changes and expect those changes to be
# present. For example, `disable_prepared_statements` expects the
# configuration settings to always be up to date.
#
# See the following for more information:
#
# - https://gitlab.com/gitlab-org/release/retrospectives/-/issues/39
# - https://gitlab.com/gitlab-com/gl-infra/production/-/issues/5238
scope.connection_db_config.configuration_hash.with_indifferent_access
end
def pool_size
config[:pool] || Database.default_pool_size
end
def username
config[:username] || ENV['USER']
end
def database_name
config[:database]
end
def adapter_name
config[:adapter]
end
def human_adapter_name
if postgresql?
'PostgreSQL'
else
'Unknown'
end
end
def postgresql?
adapter_name.casecmp('postgresql') == 0
end
def db_config_with_default_pool_size
db_config_object = scope.connection_db_config
config = db_config_object
.configuration_hash
.merge(pool: Database.default_pool_size)
ActiveRecord::DatabaseConfigurations::HashConfig.new(
db_config_object.env_name,
db_config_object.name,
config
)
end
# Disables prepared statements for the current database connection.
def disable_prepared_statements
db_config_object = scope.connection_db_config
config = db_config_object.configuration_hash.merge(prepared_statements: false)
hash_config = ActiveRecord::DatabaseConfigurations::HashConfig.new(
db_config_object.env_name,
db_config_object.name,
config
)
scope.establish_connection(hash_config)
end
# Check whether the underlying database is in read-only mode
def db_read_only?
pg_is_in_recovery =
scope
.connection
.execute('SELECT pg_is_in_recovery()')
.first
.fetch('pg_is_in_recovery')
Gitlab::Utils.to_boolean(pg_is_in_recovery)
end
def db_read_write?
!db_read_only?
end
def version
@version ||= database_version.match(/\A(?:PostgreSQL |)([^\s]+).*\z/)[1]
end
def database_version
connection.execute("SELECT VERSION()").first['version']
end
def postgresql_minimum_supported_version?
version.to_f >= MINIMUM_POSTGRES_VERSION
end
# Bulk inserts a number of rows into a table, optionally returning their
# IDs.
#
# table - The name of the table to insert the rows into.
# rows - An Array of Hash instances, each mapping the columns to their
# values.
# return_ids - When set to true the return value will be an Array of IDs of
# the inserted rows
# disable_quote - A key or an Array of keys to exclude from quoting (You
# become responsible for protection from SQL injection for
# these keys!)
# on_conflict - Defines an upsert. Values can be: :disabled (default) or
# :do_nothing
def bulk_insert(table, rows, return_ids: false, disable_quote: [], on_conflict: nil)
return if rows.empty?
keys = rows.first.keys
columns = keys.map { |key| connection.quote_column_name(key) }
disable_quote = Array(disable_quote).to_set
tuples = rows.map do |row|
keys.map do |k|
disable_quote.include?(k) ? row[k] : connection.quote(row[k])
end
end
sql = <<-EOF
INSERT INTO #{table} (#{columns.join(', ')})
VALUES #{tuples.map { |tuple| "(#{tuple.join(', ')})" }.join(', ')}
EOF
sql = "#{sql} ON CONFLICT DO NOTHING" if on_conflict == :do_nothing
sql = "#{sql} RETURNING id" if return_ids
result = connection.execute(sql)
if return_ids
result.values.map { |tuple| tuple[0].to_i }
else
[]
end
end
def cached_column_exists?(table_name, column_name)
connection
.schema_cache.columns_hash(table_name)
.has_key?(column_name.to_s)
end
def cached_table_exists?(table_name)
exists? && connection.schema_cache.data_source_exists?(table_name)
end
def exists?
# We can't _just_ check if `connection` raises an error, as it will
# point to a `ConnectionProxy`, and obtaining those doesn't involve any
# database queries. So instead we obtain the database version, which is
# cached after the first call.
connection.schema_cache.database_version
true
rescue StandardError
false
end
def system_id
row = connection
.execute('SELECT system_identifier FROM pg_control_system()')
.first
row['system_identifier']
end
def pg_wal_lsn_diff(location1, location2)
lsn1 = connection.quote(location1)
lsn2 = connection.quote(location2)
query = <<-SQL.squish
SELECT pg_wal_lsn_diff(#{lsn1}, #{lsn2})
AS result
SQL
row = connection.select_all(query).first
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
def inside_transaction?
base = Rails.env.test? ? @open_transactions_baseline : 0
scope.connection.open_transactions > base
end
# These methods that access @open_transactions_baseline are not
# thread-safe. These are fine though because we only call these in
# RSpec's main thread. If we decide to run specs multi-threaded, we would
# need to use something like ThreadGroup to keep track of this value
def set_open_transactions_baseline
@open_transactions_baseline = scope.connection.open_transactions
end
def reset_open_transactions_baseline
@open_transactions_baseline = 0
end
private
def connection
scope.connection
end
end
end
end
Gitlab::Database::Connection.prepend_mod_with('Gitlab::Database::Connection')
|