Apply MR suggestions

- Fix specs
- Don't initialize db metrics for logger
- Dont log db count in lograge if there is no db_duration_s

config/initializers/lograge.rb

@@ -15,6 +15,7 @@ unless Gitlab::Runtime.sidekiq?
       data[:db_duration_s] = Gitlab::Utils.ms_to_round_sec(data.delete(:db)) if data[:db]
       data[:view_duration_s] = Gitlab::Utils.ms_to_round_sec(data.delete(:view)) if data[:view]
       data[:duration_s] = Gitlab::Utils.ms_to_round_sec(data.delete(:duration)) if data[:duration]
+      data.merge!(::Gitlab::Metrics::Subscribers::ActiveRecord.db_counter_payload) if data[:db_duration_s]
 
       data
     end

lib/gitlab/lograge/custom_options.rb

@@ -20,8 +20,6 @@ module Gitlab
           username: event.payload[:username],
           ua: event.payload[:ua]
         }
-        payload.merge!(::Gitlab::Metrics::Subscribers::ActiveRecord.db_counter_payload)
-
         payload.merge!(event.payload[:metadata]) if event.payload[:metadata]
 
         ::Gitlab::InstrumentationHelper.add_instrumentation_data(payload)

lib/gitlab/metrics/subscribers/active_record.rb

@@ -24,15 +24,11 @@ module Gitlab
         end
 
         def self.db_counter_payload
-          payload = {}
+          return {} unless Gitlab::SafeRequestStore.active?
 
-          if Gitlab::SafeRequestStore.active?
-            DB_COUNTERS.each do |counter|
-              payload[counter] = Gitlab::SafeRequestStore[counter] if Gitlab::SafeRequestStore.exist?(counter)
-            end
-          end
-
-          payload
+          DB_COUNTERS.map do |counter|
+            [counter, Gitlab::SafeRequestStore[counter].to_i]
+          end.to_h
         end
 
         private
@@ -48,7 +44,6 @@ module Gitlab
         end
 
         def increment_db_counters(payload)
-          initialize_cache_keys
           increment(:db_count)
 
           if payload.fetch(:cached, payload[:name] == 'CACHE')
@@ -58,19 +53,11 @@ module Gitlab
           increment(:db_write_count) unless select_sql_command?(payload)
         end
 
-        def initialize_cache_keys
-          return unless Gitlab::SafeRequestStore.active?
-
-          DB_COUNTERS.each do |counter|
-            Gitlab::SafeRequestStore[counter] = 0 unless Gitlab::SafeRequestStore.exist?(counter)
-          end
-        end
-
         def increment(counter)
           current_transaction.increment(counter, 1)
 
           if Gitlab::SafeRequestStore.active?
-            Gitlab::SafeRequestStore[counter] += 1
+            Gitlab::SafeRequestStore[counter] = Gitlab::SafeRequestStore[counter].to_i + 1
           end
         end
 

spec/initializers/lograge_spec.rb

@@ -152,5 +152,46 @@ RSpec.describe 'lograge', type: :request do
         expect(log_data['etag_route']).to eq(etag_route)
       end
     end
+
+    context 'with transaction' do
+      let(:transaction) { Gitlab::Metrics::WebTransaction.new({}) }
+
+      before do
+        allow(Gitlab::Metrics::Transaction).to receive(:current).and_return(transaction)
+      end
+
+      context 'when RequestStore is enabled', :request_store do
+        context 'with db payload' do
+          it 'includes db counters', :request_store do
+            ActiveRecord::Base.connection.execute('SELECT pg_sleep(0.1);')
+            subscriber.process_action(event)
+
+            expect(log_data).to include("db_count" => 1, "db_write_count" => 0, "db_cached_count" => 0)
+          end
+        end
+
+        context 'with db payload' do
+          before do
+            event.payload.except!(:db_runtime)
+          end
+          it 'does not include db counters', :request_store do
+            subscriber.process_action(event)
+
+            expect(log_data).not_to include("db_count" => 0, "db_write_count" => 0, "db_cached_count" => 0)
+          end
+        end
+      end
+
+      context 'when RequestStore is disabled' do
+        context 'with db payload' do
+          it 'includes db counters' do
+            ActiveRecord::Base.connection.execute('SELECT pg_sleep(0.1);')
+            subscriber.process_action(event)
+
+            expect(log_data).not_to include("db_count" => 1, "db_write_count" => 0, "db_cached_count" => 0)
+          end
+        end
+      end
+    end
   end
 end

spec/lib/gitlab/lograge/custom_options_spec.rb

@@ -44,20 +44,6 @@ RSpec.describe Gitlab::Lograge::CustomOptions do
       end
     end
 
-    context 'with transaction' do
-      let(:transaction) { Gitlab::Metrics::WebTransaction.new({}) }
-
-      before do
-        allow(Gitlab::Metrics::Transaction).to receive(:current).and_return(transaction)
-      end
-
-      it 'adds db counters', :request_store do
-        ActiveRecord::Base.connection.execute('SELECT pg_sleep(0.1);')
-
-        expect(subject).to include(db_count: 1, db_write_count: 0, db_cached_count: 0)
-      end
-    end
-
     it 'adds the user id' do
       expect(subject[:user_id]).to eq('test')
     end

spec/lib/gitlab/metrics/subscribers/active_record_spec.rb

@@ -52,7 +52,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
             subscriber.sql(event)
 
             described_class::DB_COUNTERS.each do |counter|
-              expect(Gitlab::SafeRequestStore[counter]).to eq expected_counters[counter]
+              expect(Gitlab::SafeRequestStore[counter].to_i).to eq expected_counters[counter]
             end
           end
 
@@ -62,7 +62,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
                 subscriber.sql(event)
 
                 described_class::DB_COUNTERS.each do |counter|
-                  expect(Gitlab::SafeRequestStore[counter]).to eq expected_counters[counter]
+                  expect(Gitlab::SafeRequestStore[counter].to_i).to eq expected_counters[counter]
                 end
               end
             end
@@ -256,9 +256,15 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
       end
 
       context 'when query is not executed' do
-        let(:expected_payload) { {} }
+        let(:expected_payload) do
+          {
+            db_count: 0,
+            db_cached_count: 0,
+            db_write_count: 0
+          }
+        end
 
-        it 'returns empty payload' do
+        it 'returns correct payload' do
           expect(described_class.db_counter_payload).to eq(expected_payload)
         end
       end
@@ -267,7 +273,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
     context 'when RequestStore is disabled' do
       let(:expected_payload) { {} }
 
-      it 'returns correct payload' do
+      it 'returns empty payload' do
         subscriber.sql(event)
 
         expect(described_class.db_counter_payload).to eq(expected_payload)