Extend `marginalia` to provide `db_config_name`

The `db_config_name` describes a type of connection being
used by a given query.

Changelog: added

config/initializers/0_marginalia.rb

@@ -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

lib/gitlab/database.rb

@@ -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
+    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
 
-      'unknown'
+    # 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?

lib/gitlab/database/load_balancing.rb

@@ -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

lib/gitlab/marginalia/comment.rb

@@ -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

spec/lib/gitlab/database/schema_migrations/context_spec.rb

@@ -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

spec/lib/gitlab/database_spec.rb

@@ -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'"

spec/lib/marginalia_spec.rb

@@ -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
 
@@ -100,9 +125,10 @@ RSpec.describe 'Marginalia spec' do
 
       let(:component_map) do
         {
-          "application" => "sidekiq",
-          "endpoint_id" => "ActionMailer::MailDeliveryJob",
-          "jid"         => delivery_job.job_id
+          "application"    => "sidekiq",
+          "endpoint_id"    => "ActionMailer::MailDeliveryJob",
+          "jid"            => delivery_job.job_id,
+          "db_config_name" => "main"
         }
       end
 

spec/spec_helper.rb

@@ -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

spec/support/database/multiple_databases.rb

+# 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