Collect database host metrics for all primaries

Update the DatabaseSampler to collect metrics for all primaries defined
in Gitlab::Database.database_base_models, rather than ActiveRecord::Base
only.

lib/gitlab/metrics/samplers/database_sampler.rb

@@ -38,16 +38,19 @@ module Gitlab
         end
 
         def host_stats
-          return [] unless ActiveRecord::Base.connected?
+          Gitlab::Database.database_base_models.each_value.with_object([]) do |base_model, stats|
+            next unless base_model.connected?
 
-          [{ labels: labels_for_class(ActiveRecord::Base), stats: ActiveRecord::Base.connection_pool.stat }]
+            stats << { labels: labels_for_class(base_model), stats: base_model.connection_pool.stat }
+          end
         end
 
         def labels_for_class(klass)
           {
             host: klass.connection_db_config.host,
             port: klass.connection_db_config.configuration_hash[:port],
-            class: klass.to_s
+            class: klass.to_s,
+            db_config_name: klass.connection_db_config.name
           }
         end
       end

spec/lib/gitlab/metrics/samplers/database_sampler_spec.rb

@@ -8,44 +8,62 @@ RSpec.describe Gitlab::Metrics::Samplers::DatabaseSampler do
   it_behaves_like 'metrics sampler', 'DATABASE_SAMPLER'
 
   describe '#sample' do
+    let(:active_record_labels) do
+      {
+        class: 'ActiveRecord::Base',
+        host: ApplicationRecord.database.config['host'],
+        port: ApplicationRecord.database.config['port'],
+        db_config_name: 'main'
+      }
+    end
+
+    let(:ci_application_record_labels) do
+      {
+        class: 'Ci::ApplicationRecord',
+        host: Ci::ApplicationRecord.database.config['host'],
+        port: Ci::ApplicationRecord.database.config['port'],
+        db_config_name: 'ci'
+      }
+    end
+
     before do
       described_class::METRIC_DESCRIPTIONS.each_key do |metric|
         allow(subject.metrics[metric]).to receive(:set)
       end
-    end
 
-    context 'for ActiveRecord::Base' do
-      let(:labels) do
-        {
-          class: 'ActiveRecord::Base',
-          host: ApplicationRecord.database.config['host'],
-          port: ApplicationRecord.database.config['port']
-        }
+      allow(Gitlab::Database).to receive(:database_base_models)
+        .and_return({ main: ActiveRecord::Base, ci: Ci::ApplicationRecord })
     end
 
-      context 'when the database is connected' do
+    context 'when the database is connected', :add_ci_connection do
       it 'samples connection pool statistics' do
-          expect(subject.metrics[:size]).to receive(:set).with(labels, a_value >= 1)
-          expect(subject.metrics[:connections]).to receive(:set).with(labels, a_value >= 1)
-          expect(subject.metrics[:busy]).to receive(:set).with(labels, a_value >= 1)
-          expect(subject.metrics[:dead]).to receive(:set).with(labels, a_value >= 0)
-          expect(subject.metrics[:waiting]).to receive(:set).with(labels, a_value >= 0)
+        expect(subject.metrics[:size]).to receive(:set).with(active_record_labels, a_value >= 1)
+        expect(subject.metrics[:connections]).to receive(:set).with(active_record_labels, a_value >= 1)
+        expect(subject.metrics[:busy]).to receive(:set).with(active_record_labels, a_value >= 1)
+        expect(subject.metrics[:dead]).to receive(:set).with(active_record_labels, a_value >= 0)
+        expect(subject.metrics[:waiting]).to receive(:set).with(active_record_labels, a_value >= 0)
+
+        expect(subject.metrics[:size]).to receive(:set).with(ci_application_record_labels, a_value >= 1)
+        expect(subject.metrics[:connections]).to receive(:set).with(ci_application_record_labels, a_value >= 1)
+        expect(subject.metrics[:busy]).to receive(:set).with(ci_application_record_labels, a_value >= 1)
+        expect(subject.metrics[:dead]).to receive(:set).with(ci_application_record_labels, a_value >= 0)
+        expect(subject.metrics[:waiting]).to receive(:set).with(ci_application_record_labels, a_value >= 0)
 
         subject.sample
       end
     end
 
-      context 'when the database is not connected' do
+    context 'when a database is not connected', :add_ci_connection do
       before do
-          allow(ActiveRecord::Base).to receive(:connected?).and_return(false)
+        allow(Ci::ApplicationRecord).to receive(:connected?).and_return(false)
       end
 
-        it 'records no samples' do
-          expect(subject.metrics[:size]).not_to receive(:set).with(labels, anything)
+      it 'records no samples for that database' do
+        expect(subject.metrics[:size]).to receive(:set).with(active_record_labels, anything)
+        expect(subject.metrics[:size]).not_to receive(:set).with(ci_application_record_labels, anything)
 
         subject.sample
       end
     end
   end
-  end
 end