Track db calls in metrics

- Add db calls to Metrics::Subscribers::ActiveRecord
- Add prometheus metrics for rails
- Add prometheus metrics for sidekiq
- Add metrics to lograge
- Add metrics to StructuredLogger
- Add missing specs

changelogs/unreleased/216835-instrument-db-calls.yml

+---
+title: Add number of database calls to Prometheus metrics and logs for sidekiq and
+  request
+merge_request: 32131
+author:
+type: added

doc/administration/monitoring/prometheus/gitlab_metrics.md

@@ -92,6 +92,9 @@ The following metrics are available:
 | `gitlab_view_rendering_duration_seconds`                       | Histogram |                   10.2 | Duration for views (histogram)                                                                      | `controller`, `action`, `view`                      |
 | `http_requests_total`                                          | Counter   |                    9.4 | Rack request count                                                                                  | `method`                                            |
 | `http_request_duration_seconds`                                | Histogram |                    9.4 | HTTP response time from rack middleware                                                             | `method`, `status`                                  |
+| `gitlab_transaction_db_count_total`                            | Counter   |                   13.1 | Counter for total number of sql calls                                                               | `controller`, `action`                              |
+| `gitlab_transaction_db_write_count_total`                      | Counter   |                   13.1 | Counter for total number of write sql calls                                                         | `controller`, `action`                              |
+| `gitlab_transaction_db_cached_count_total`                     | Counter   |                   13.1 | Counter for total number of cached sql calls                                                        | `controller`, `action`                              |
 | `http_redis_requests_duration_seconds`                         | Histogram |                   13.1 | Redis requests duration during web transactions                                                     | `controller`, `action`                              |
 | `http_redis_requests_total`                                    | Counter   |                   13.1 | Redis requests count during web transactions                                                        | `controller`, `action`                              |
 | `http_elasticsearch_requests_duration_seconds` **(STARTER)**   | Histogram |                   13.1 | Elasticsearch requests duration during web transactions                                             | `controller`, `action`                              |

lib/gitlab/lograge/custom_options.rb

@@ -12,7 +12,6 @@ module Gitlab
         params = event
           .payload[:params]
           .each_with_object([]) { |(k, v), array| array << { key: k, value: v } unless IGNORE_PARAMS.include?(k) }
-
         payload = {
           time: Time.now.utc.iso8601(3),
           params: Gitlab::Utils::LogLimitedArray.log_limited_array(params, sentinel: LIMITED_ARRAY_SENTINEL),
@@ -21,6 +20,7 @@ module Gitlab
           username: event.payload[:username],
           ua: event.payload[:ua]
         }
+        add_db_counters!(payload)
 
         payload.merge!(event.payload[:metadata]) if event.payload[:metadata]
 
@@ -46,6 +46,16 @@ module Gitlab
 
         payload
       end
+
+      def self.add_db_counters!(payload)
+        current_transaction = Gitlab::Metrics::Transaction.current
+        if current_transaction
+          payload[:db_count] = current_transaction.get(:db_count, :counter).to_i
+          payload[:db_write_count] = current_transaction.get(:db_write_count, :counter).to_i
+          payload[:db_cached_count] = current_transaction.get(:db_cached_count, :counter).to_i
+        end
+      end
+      private_class_method :add_db_counters!
     end
   end
 end

lib/gitlab/metrics/sidekiq_middleware.rb

@@ -6,19 +6,30 @@ module Gitlab
     #
     # This middleware is intended to be used as a server-side middleware.
     class SidekiqMiddleware
-      def call(worker, message, queue)
+      def call(worker, payload, queue)
         trans = BackgroundTransaction.new(worker.class)
 
         begin
           # Old gitlad-shell messages don't provide enqueued_at/created_at attributes
-          trans.set(:sidekiq_queue_duration, Time.now.to_f - (message['enqueued_at'] || message['created_at'] || 0))
+          enqueued_at = payload['enqueued_at'] || payload['created_at'] || 0
+          trans.set(:sidekiq_queue_duration, Time.current.to_f - enqueued_at)
           trans.run { yield }
         rescue Exception => error # rubocop: disable Lint/RescueException
           trans.add_event(:sidekiq_exception)
 
           raise error
+        ensure
+          add_info_to_payload(payload, trans)
         end
       end
+
+      private
+
+      def add_info_to_payload(payload, trans)
+        payload[:db_count] = trans.get(:db_count, :counter).to_i
+        payload[:db_write_count] = trans.get(:db_write_count, :counter).to_i
+        payload[:db_cached_count] = trans.get(:db_cached_count, :counter).to_i
+      end
     end
   end
 end

lib/gitlab/metrics/subscribers/active_record.rb

@@ -9,6 +9,7 @@ module Gitlab
         attach_to :active_record
 
         IGNORABLE_SQL = %w{BEGIN COMMIT}.freeze
+        DB_COUNTERS = %i{db_count db_write_count db_cached_count}.freeze
 
         def sql(event)
           return unless current_transaction
@@ -19,8 +20,7 @@ module Gitlab
 
           self.class.gitlab_sql_duration_seconds.observe(current_transaction.labels, event.duration / 1000.0)
 
-          current_transaction.increment(:sql_duration, event.duration, false)
-          current_transaction.increment(:sql_count, 1, false)
+          increment_db_counters(payload)
         end
 
         private
@@ -31,6 +31,20 @@ module Gitlab
           buckets [0.01, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0]
         end
 
+        def select_sql_command?(payload)
+          payload[:sql].match(/\A((?!(.*[^\w'"](DELETE|UPDATE|INSERT INTO)[^\w'"])))(WITH.*)?(SELECT)((?!(FOR UPDATE|FOR SHARE)).)*$/i)
+        end
+
+        def increment_db_counters(payload)
+          current_transaction.increment(:db_count, 1)
+
+          if payload.fetch(:cached, payload[:name] == 'CACHE')
+            current_transaction.increment(:db_cached_count, 1)
+          end
+
+          current_transaction.increment(:db_write_count, 1) unless select_sql_command?(payload)
+        end
+
         def current_transaction
           Transaction.current
         end

lib/gitlab/metrics/transaction.rb

@@ -16,7 +16,7 @@ module Gitlab
       # The series to store events (e.g. Git pushes) in.
       EVENT_SERIES = 'events'
 
-      attr_reader :tags, :method
+      attr_reader :method
 
       def self.current
         Thread.current[THREAD_KEY]
@@ -28,8 +28,6 @@ module Gitlab
         @started_at = nil
         @finished_at = nil
 
-        @tags = {}
-
         @memory_before = 0
         @memory_after = 0
       end
@@ -94,6 +92,12 @@ module Gitlab
         self.class.transaction_metric(name, :gauge).set(labels, value) if use_prometheus
       end
 
+      def get(name, type, tags = {})
+        metric = self.class.transaction_metric(name, type)
+
+        metric.get(filter_tags(tags).merge(labels))
+      end
+
       def labels
         BASE_LABELS
       end

lib/gitlab/sidekiq_logging/structured_logger.rb

@@ -36,6 +36,10 @@ module Gitlab
         )
       end
 
+      def add_db_counters!(job, output_payload)
+        output_payload.merge!(job.slice(*::Gitlab::Metrics::Subscribers::ActiveRecord::DB_COUNTERS))
+      end
+
       def log_job_start(payload)
         payload['message'] = "#{base_message(payload)}: start"
         payload['job_status'] = 'start'
@@ -50,6 +54,7 @@ module Gitlab
         payload = payload.dup
         add_instrumentation_keys!(job, payload)
         add_logging_extras!(job, payload)
+        add_db_counters!(job, payload)
 
         elapsed_time = elapsed(started_time)
         add_time_keys!(elapsed_time, payload)

spec/lib/gitlab/lograge/custom_options_spec.rb

@@ -44,6 +44,18 @@ 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' do
+        expect(subject).to include(:db_count, :db_write_count, :db_cached_count)
+      end
+    end
+
     it 'adds the user id' do
       expect(subject[:user_id]).to eq('test')
     end

spec/lib/gitlab/metrics/sidekiq_middleware_spec.rb

@@ -10,17 +10,19 @@ describe Gitlab::Metrics::SidekiqMiddleware do
     it 'tracks the transaction' do
       worker = double(:worker, class: double(:class, name: 'TestWorker'))
 
-      expect(Gitlab::Metrics::BackgroundTransaction).to receive(:new)
-        .with(worker.class)
-        .and_call_original
+      expect_next_instance_of(Gitlab::Metrics::BackgroundTransaction) do |transaction|
+        expect(transaction).to receive(:set).with(:sidekiq_queue_duration, instance_of(Float))
+        expect(transaction).to receive(:increment).with(:db_count, 1)
+      end
 
-      expect_any_instance_of(Gitlab::Metrics::Transaction).to receive(:set)
-        .with(:sidekiq_queue_duration, instance_of(Float))
+      middleware.call(worker, message, :test) do
+        ActiveRecord::Base.connection.execute('SELECT pg_sleep(0.1);')
+      end
 
-      middleware.call(worker, message, :test) { nil }
+      expect(message).to include(:db_count, :db_write_count, :db_cached_count)
     end
 
-    it 'tracks the transaction (for messages without `enqueued_at`)' do
+    it 'tracks the transaction (for messages without `enqueued_at`)', :aggregate_failures do
       worker = double(:worker, class: double(:class, name: 'TestWorker'))
 
       expect(Gitlab::Metrics::BackgroundTransaction).to receive(:new)
@@ -33,7 +35,7 @@ describe Gitlab::Metrics::SidekiqMiddleware do
       middleware.call(worker, {}, :test) { nil }
     end
 
-    it 'tracks any raised exceptions' do
+    it 'tracks any raised exceptions', :aggregate_failures do
       worker = double(:worker, class: double(:class, name: 'TestWorker'))
 
       expect_any_instance_of(Gitlab::Metrics::Transaction)
@@ -44,6 +46,8 @@ describe Gitlab::Metrics::SidekiqMiddleware do
 
       expect { middleware.call(worker, message, :test) }
         .to raise_error(RuntimeError)
+
+      expect(message).to include(:db_count, :db_write_count, :db_cached_count)
     end
   end
 end

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

@@ -6,10 +6,15 @@ describe Gitlab::Metrics::Subscribers::ActiveRecord do
   let(:env) { {} }
   let(:transaction) { Gitlab::Metrics::WebTransaction.new(env) }
   let(:subscriber)  { described_class.new }
+  let(:payload) { { sql: 'SELECT * FROM users WHERE id = 10' } }
 
   let(:event) do
-    double(:event, duration: 2,
-                   payload:  { sql: 'SELECT * FROM users WHERE id = 10' })
+    double(
+      :event,
+      name: 'sql.active_record',
+      duration: 2,
+      payload:  payload
+    )
   end
 
   describe '#sql' do
@@ -23,28 +28,133 @@ describe Gitlab::Metrics::Subscribers::ActiveRecord do
     end
 
     describe 'with a current transaction' do
-      it 'observes sql_duration metric' do
+      shared_examples 'read only query' do
+        it 'increments only db count value' do
+          allow(subscriber).to receive(:current_transaction)
+                                  .at_least(:once)
+                                  .and_return(transaction)
+
+          expect(transaction).to receive(:increment)
+                                   .with(:db_count, 1)
+
+          expect(transaction).not_to receive(:increment)
+                                       .with(:db_cached_count, 1)
+
+          expect(transaction).not_to receive(:increment)
+                                       .with(:db_write_count, 1)
+
+          subscriber.sql(event)
+        end
+      end
+
+      shared_examples 'write query' do
+        it 'increments db_write_count and db_count value' do
           expect(subscriber).to receive(:current_transaction)
                                   .at_least(:once)
                                   .and_return(transaction)
-        expect(described_class.send(:gitlab_sql_duration_seconds)).to receive(:observe).with({}, 0.002)
+
+          expect(transaction).to receive(:increment)
+                                   .with(:db_count, 1)
+
+          expect(transaction).not_to receive(:increment)
+                                       .with(:db_cached_count, 1)
+
+          expect(transaction).to receive(:increment)
+                                       .with(:db_write_count, 1)
+
           subscriber.sql(event)
         end
+      end
 
-      it 'increments the :sql_duration value' do
+      shared_examples 'cached query' do
+        it 'increments db_cached_count and db_count value' do
           expect(subscriber).to receive(:current_transaction)
                                   .at_least(:once)
                                   .and_return(transaction)
 
           expect(transaction).to receive(:increment)
-          .with(:sql_duration, 2, false)
+                                   .with(:db_count, 1)
 
           expect(transaction).to receive(:increment)
-          .with(:sql_count, 1, false)
+                                       .with(:db_cached_count, 1)
+
+          expect(transaction).not_to receive(:increment)
+                                   .with(:db_write_count, 1)
+
+          subscriber.sql(event)
+        end
+      end
 
+      it 'observes sql_duration metric' do
+        expect(subscriber).to receive(:current_transaction)
+                                .at_least(:once)
+                                .and_return(transaction)
+        expect(described_class.send(:gitlab_sql_duration_seconds)).to receive(:observe).with({}, 0.002)
         subscriber.sql(event)
       end
 
+      it_behaves_like 'read only query'
+
+      context 'with select for update sql event' do
+        let(:payload) { { sql: 'SELECT * FROM users WHERE id = 10 FOR UPDATE' } }
+
+        it_behaves_like 'write query'
+      end
+
+      context 'with common table expression' do
+        context 'with insert' do
+          let(:payload) { { sql: 'WITH archived_rows AS (SELECT * FROM users WHERE archived = true) INSERT INTO products_log SELECT * FROM archived_rows' } }
+
+          it_behaves_like 'write query'
+        end
+
+        context 'with only select' do
+          let(:payload) { { sql: 'WITH active_milestones AS (SELECT COUNT(*), state FROM milestones GROUP BY state) SELECT * FROM active_milestones' } }
+
+          it_behaves_like 'read only query'
+        end
+      end
+
+      context 'with delete sql event' do
+        let(:payload) { { sql: 'DELETE FROM users where id = 10' } }
+
+        it_behaves_like 'write query'
+      end
+
+      context 'with insert sql event' do
+        let(:payload) { { sql: 'INSERT INTO project_ci_cd_settings (project_id) SELECT id FROM projects' } }
+
+        it_behaves_like 'write query'
+      end
+
+      context 'with update sql event' do
+        let(:payload) { { sql: 'UPDATE users SET admin = true WHERE id = 10' } }
+
+        it_behaves_like 'write query'
+      end
+
+      context 'with cached payload ' do
+        let(:payload) do
+          {
+            sql: 'SELECT * FROM users WHERE id = 10',
+            cached: true
+          }
+        end
+
+        it_behaves_like 'cached query'
+      end
+
+      context 'with cached payload name' do
+        let(:payload) do
+          {
+           sql: 'SELECT * FROM users WHERE id = 10',
+           name: 'CACHE'
+          }
+        end
+
+        it_behaves_like 'cached query'
+      end
+
       context 'events are internal to Rails or irrelevant' do
         let(:schema_event) do
           double(

spec/lib/gitlab/metrics/transaction_spec.rb

@@ -65,18 +65,16 @@ describe Gitlab::Metrics::Transaction do
   describe '#add_event' do
     let(:prometheus_metric) { instance_double(Prometheus::Client::Counter, increment: nil) }
 
-    before do
-      allow(described_class).to receive(:transaction_metric).and_return(prometheus_metric)
-    end
-
     it 'adds a metric' do
       expect(prometheus_metric).to receive(:increment)
+      expect(described_class).to receive(:fetch_metric).with(:counter, :gitlab_transaction_event_meow_total).and_return(prometheus_metric)
 
       transaction.add_event(:meow)
     end
 
     it 'allows tracking of custom tags' do
       expect(prometheus_metric).to receive(:increment).with(hash_including(animal: "dog"))
+      expect(described_class).to receive(:fetch_metric).with(:counter, :gitlab_transaction_event_bau_total).and_return(prometheus_metric)
 
       transaction.add_event(:bau, animal: 'dog')
     end
@@ -84,6 +82,7 @@ describe Gitlab::Metrics::Transaction do
     context 'with sensitive tags' do
       before do
         transaction.add_event(:baubau, **sensitive_tags.merge(sane: 'yes'))
+        allow(described_class).to receive(:transaction_metric).and_return(prometheus_metric)
       end
 
       it 'filters tags' do
@@ -93,4 +92,37 @@ describe Gitlab::Metrics::Transaction do
       end
     end
   end
+
+  describe '#increment' do
+    let(:prometheus_metric) { instance_double(Prometheus::Client::Counter, increment: nil) }
+
+    it 'adds a metric' do
+      expect(prometheus_metric).to receive(:increment).with(hash_including(:action, :controller), 1)
+      expect(described_class).to receive(:fetch_metric).with(:counter, :gitlab_transaction_meow_total).and_return(prometheus_metric)
+
+      transaction.increment(:meow, 1)
+    end
+  end
+
+  describe '#set' do
+    let(:prometheus_metric) { instance_double(Prometheus::Client::Gauge, set: nil) }
+
+    it 'adds a metric' do
+      expect(prometheus_metric).to receive(:set).with(hash_including(:action, :controller), 1)
+      expect(described_class).to receive(:fetch_metric).with(:gauge, :gitlab_transaction_meow_total).and_return(prometheus_metric)
+
+      transaction.set(:meow, 1)
+    end
+  end
+
+  describe '#get' do
+    let(:prometheus_metric) { instance_double(Prometheus::Client::Counter, get: nil) }
+
+    it 'gets a metric' do
+      expect(described_class).to receive(:fetch_metric).with(:counter, :gitlab_transaction_meow_total).and_return(prometheus_metric)
+      expect(prometheus_metric).to receive(:get)
+
+      transaction.get(:meow, :counter)
+    end
+  end
 end

spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb

@@ -21,7 +21,10 @@ describe Gitlab::SidekiqLogging::StructuredLogger do
         "correlation_id" => 'cid',
         "error_message" => "wrong number of arguments (2 for 3)",
         "error_class" => "ArgumentError",
-        "error_backtrace" => []
+        "error_backtrace" => [],
+        "db_count" => 1,
+        "db_write_count" => 0,
+        "db_cached_count" => 0
       }
     end
 
@@ -197,7 +200,10 @@ describe Gitlab::SidekiqLogging::StructuredLogger do
 
       let(:expected_end_payload_with_db) do
         expected_end_payload.merge(
-          'db_duration_s' => a_value >= 0.1
+          'db_duration_s' => a_value >= 0.1,
+          'db_count' => 1,
+          'db_cached_count' => 0,
+          'db_write_count' => 0
         )
       end