debian-mirror-gitlab/lib/gitlab/database/connection.rb
2021-11-11 11:23:49 +05:30

260 lines
7.9 KiB
Ruby

# 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')