Commit 759bb360 authored by Adam Hegyi's avatar Adam Hegyi

Merge branch 'many-dbs-margnalia-comment' into 'master'

Extend `marginalia` to provide `db_config_name`

See merge request gitlab-org/gitlab!67328
parents 41e7f385 24e07a2a
......@@ -13,7 +13,7 @@ require 'marginalia'
# matching against the raw SQL, and prepending the comment prevents color
# coding from working in the development log.
Marginalia::Comment.prepend_comment = true if Rails.env.production?
Marginalia::Comment.components = [:application, :correlation_id, :jid, :endpoint_id]
Marginalia::Comment.components = [:application, :correlation_id, :jid, :endpoint_id, :db_config_name]
# As mentioned in https://github.com/basecamp/marginalia/pull/93/files,
# adding :line has some overhead because a regexp on the backtrace has
......
......@@ -198,14 +198,30 @@ module Gitlab
::ActiveRecord::Base.configurations.configs_for(env_name: Rails.env).map(&:name)
end
def self.db_config_name(ar_connection)
if ar_connection.respond_to?(:pool) &&
ar_connection.pool.respond_to?(:db_config) &&
ar_connection.pool.db_config.respond_to?(:name)
return ar_connection.pool.db_config.name
end
'unknown'
def self.db_config_for_connection(connection)
return unless connection
# The LB connection proxy does not have a direct db_config
# that can be referenced
return if connection.is_a?(::Gitlab::Database::LoadBalancing::ConnectionProxy)
# During application init we might receive `NullPool`
return unless connection.respond_to?(:pool) &&
connection.pool.respond_to?(:db_config)
connection.pool.db_config
end
# At the moment, the connection can only be retrieved by
# Gitlab::Database::LoadBalancer#read or #read_write or from the
# ActiveRecord directly. Therefore, if the load balancer doesn't
# recognize the connection, this method returns the primary role
# directly. In future, we may need to check for other sources.
# Expected returned names:
# main, main_replica, ci, ci_replica, unknown
def self.db_config_name(connection)
db_config = db_config_for_connection(connection)
db_config&.name || 'unknown'
end
def self.read_only?
......
......@@ -79,24 +79,12 @@ module Gitlab
].freeze
# Returns the role (primary/replica) of the database the connection is
# connecting to. At the moment, the connection can only be retrieved by
# Gitlab::Database::LoadBalancer#read or #read_write or from the
# ActiveRecord directly. Therefore, if the load balancer doesn't
# recognize the connection, this method returns the primary role
# directly. In future, we may need to check for other sources.
# connecting to.
def self.db_role_for_connection(connection)
return ROLE_UNKNOWN unless connection
db_config = Database.db_config_for_connection(connection)
return ROLE_UNKNOWN unless db_config
# The connection proxy does not have a role assigned
# as this is dependent on a execution context
return ROLE_UNKNOWN if connection.is_a?(ConnectionProxy)
# During application init we might receive `NullPool`
return ROLE_UNKNOWN unless connection.respond_to?(:pool) &&
connection.pool.respond_to?(:db_config) &&
connection.pool.db_config.respond_to?(:name)
if connection.pool.db_config.name.ends_with?(LoadBalancer::REPLICA_SUFFIX)
if db_config.name.ends_with?(LoadBalancer::REPLICA_SUFFIX)
ROLE_REPLICA
else
ROLE_PRIMARY
......
......@@ -41,6 +41,10 @@ module Gitlab
def endpoint_id
Labkit::Context.current&.get_attribute(:caller_id)
end
def db_config_name
::Gitlab::Database.db_config_name(marginalia_adapter)
end
end
end
end
......@@ -124,8 +124,4 @@ RSpec.describe Gitlab::Database::SchemaMigrations::Context do
end
end
end
def skip_if_multiple_databases_not_setup
skip 'Skipping because multiple databases not set up' unless Gitlab::Database.has_config?(:ci)
end
end
......@@ -155,23 +155,43 @@ RSpec.describe Gitlab::Database do
it { expect(described_class.nulls_first_order('column', 'DESC')).to eq 'column DESC NULLS FIRST'}
end
describe '.db_config_name' do
it 'returns the db_config name for the connection' do
connection = ActiveRecord::Base.connection
describe '.db_config_for_connection' do
context 'when the regular connection is used' do
it 'returns db_config' do
connection = ActiveRecord::Base.retrieve_connection
expect(described_class.db_config_name(connection)).to be_a(String)
expect(described_class.db_config_name(connection)).to eq(connection.pool.db_config.name)
expect(described_class.db_config_for_connection(connection)).to eq(connection.pool.db_config)
end
end
context 'when the connection is LoadBalancing::ConnectionProxy' do
it 'returns nil' do
lb_config = ::Gitlab::Database::LoadBalancing::Configuration.new(ActiveRecord::Base)
lb = ::Gitlab::Database::LoadBalancing::LoadBalancer.new(lb_config)
proxy = ::Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb)
expect(described_class.db_config_for_connection(proxy)).to be_nil
end
end
context 'when the pool is a NullPool' do
it 'returns unknown' do
it 'returns nil' do
connection = double(:active_record_connection, pool: ActiveRecord::ConnectionAdapters::NullPool.new)
expect(described_class.db_config_name(connection)).to eq('unknown')
expect(described_class.db_config_for_connection(connection)).to be_nil
end
end
end
describe '.db_config_name' do
it 'returns the db_config name for the connection' do
connection = ActiveRecord::Base.connection
expect(described_class.db_config_name(connection)).to be_a(String)
expect(described_class.db_config_name(connection)).to eq(connection.pool.db_config.name)
end
end
describe '#true_value' do
it 'returns correct value' do
expect(described_class.true_value).to eq "'t'"
......
......@@ -42,7 +42,8 @@ RSpec.describe 'Marginalia spec' do
{
"application" => "test",
"endpoint_id" => "MarginaliaTestController#first_user",
"correlation_id" => correlation_id
"correlation_id" => correlation_id,
"db_config_name" => "main"
}
end
......@@ -51,6 +52,29 @@ RSpec.describe 'Marginalia spec' do
expect(recorded.log.last).to include("#{component}:#{value}")
end
end
context 'when using CI database' do
let(:component_map) do
{
"application" => "test",
"endpoint_id" => "MarginaliaTestController#first_user",
"correlation_id" => correlation_id,
"db_config_name" => "ci"
}
end
before do |example|
skip_if_multiple_databases_not_setup
allow(User).to receive(:connection) { Ci::CiDatabaseRecord.connection }
end
it 'generates a query that includes the component and value' do
component_map.each do |component, value|
expect(recorded.log.last).to include("#{component}:#{value}")
end
end
end
end
describe 'for Sidekiq worker jobs' do
......@@ -79,7 +103,8 @@ RSpec.describe 'Marginalia spec' do
"application" => "sidekiq",
"endpoint_id" => "MarginaliaTestJob",
"correlation_id" => sidekiq_job['correlation_id'],
"jid" => sidekiq_job['jid']
"jid" => sidekiq_job['jid'],
"db_config_name" => "main"
}
end
......@@ -102,7 +127,8 @@ RSpec.describe 'Marginalia spec' do
{
"application" => "sidekiq",
"endpoint_id" => "ActionMailer::MailDeliveryJob",
"jid" => delivery_job.job_id
"jid" => delivery_job.job_id,
"db_config_name" => "main"
}
end
......
......@@ -164,6 +164,7 @@ RSpec.configure do |config|
config.include NextInstanceOf
config.include TestEnv
config.include FileReadHelpers
config.include Database::MultipleDatabases
config.include Devise::Test::ControllerHelpers, type: :controller
config.include Devise::Test::ControllerHelpers, type: :view
config.include Devise::Test::IntegrationHelpers, type: :feature
......
# frozen_string_literal: true
module Database
module MultipleDatabases
def skip_if_multiple_databases_not_setup
skip 'Skipping because multiple databases not set up' unless Gitlab::Database.has_config?(:ci)
end
end
end
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