Log DB query count per request

Author: philippthun

lib/cloud_controller/db.rb

@@ -3,6 +3,7 @@
 require 'cloud_controller/db_connection/options_factory'
 require 'cloud_controller/db_connection/finalizer'
 require 'sequel/extensions/query_length_logging'
+require 'sequel/extensions/request_query_counter'
 
 module VCAP::CloudController
   class DB
@@ -25,6 +26,7 @@ def self.connect(opts, logger)
       connection_options = VCAP::CloudController::DbConnection::OptionsFactory.build(opts)
 
       db = get_connection(opts, connection_options)
+      db.extension(:request_query_counter)
 
       if opts[:log_db_queries]
         db.logger = logger

lib/cloud_controller/logs/request_logs.rb

@@ -26,13 +26,14 @@ def start_request(request_id, env)
         @incomplete_requests.store(request_id, env)
       end
 
-      def complete_request(request_id, status, env, time_taken)
+      def complete_request(request_id, status, env, time_taken, query_count)
         return if @incomplete_requests.delete(request_id).nil?
 
         request = ActionDispatch::Request.new(env)
         @logger.info("Completed #{status} vcap-request-id: #{request_id}",
                      { status_code: status,
                        time_taken_in_ms: time_taken,
+                       db_query_count: query_count,
                        request_method: request.request_method,
                        request_fullpath: request.filtered_path })
       end

lib/sequel/extensions/request_query_counter.rb

@@ -0,0 +1,10 @@
+module Sequel
+  module RequestQueryCounter
+    def log_connection_yield(sql, conn, args=nil)
+      VCAP::Request.increment_db_query_count
+      super
+    end
+  end
+
+  Database.register_extension(:request_query_counter) { |db| db.extend(RequestQueryCounter) }
+end

lib/vcap/request.rb

@@ -55,6 +55,19 @@ def b3_span_id=(span_id)
       def b3_span_id
         Thread.current[:b3_span_id]
       end
+
+      def db_query_count
+        Thread.current[:db_query_count]
+      end
+
+      def increment_db_query_count
+        Thread.current[:db_query_count] ||= 0
+        Thread.current[:db_query_count] += 1
+      end
+
+      def reset_db_query_count
+        Thread.current[:db_query_count] = 0
+      end
     end
   end
 end

middleware/request_logs.rb

@@ -15,7 +15,8 @@ def call(env)
         # convert to milliseconds
         time_taken = (Time.now - start_time) * 1000
         time_taken = time_taken.to_i
-        @request_logs.complete_request(request_id, status, env, time_taken)
+        query_count = ::VCAP::Request.db_query_count
+        @request_logs.complete_request(request_id, status, env, time_taken, query_count)
 
         [status, headers, body]
       end

middleware/request_metrics.rb

@@ -7,6 +7,7 @@ def initialize(app, request_metrics)
       end
 
       def call(env)
+        ::VCAP::Request.reset_db_query_count
         @request_metrics.start_request
 
         status, headers, body = @app.call(env)

spec/unit/lib/cloud_controller/logs/request_logs_spec.rb

@@ -10,6 +10,7 @@ module VCAP::CloudController::Logs
     let(:fake_request) { double('request', request_method: 'request_method', ip: fake_ip, filtered_path: 'filtered_path', fullpath: fake_fullpath) }
     let(:request_id) { 'ID' }
     let(:time_taken) { 30 }
+    let(:query_count) { 5 }
     let(:env) { { 'cf.user_guid' => 'user-guid' } }
     let(:status) { 200 }
 
@@ -46,10 +47,11 @@ module VCAP::CloudController::Logs
           end
 
           it 'logs the completion of the request' do
-            request_logs.complete_request(request_id, status, env, time_taken)
+            request_logs.complete_request(request_id, status, env, time_taken, query_count)
             expect(logger).to have_received(:info).with(
               /\ACompleted 200 vcap-request-id: ID/,
               time_taken_in_ms: 30,
+              db_query_count: 5,
               request_method: 'request_method',
               request_fullpath: 'filtered_path',
               status_code: 200
@@ -59,7 +61,7 @@ module VCAP::CloudController::Logs
 
         context 'without a matching start request' do
           it 'does not log the completion of the request' do
-            request_logs.complete_request(request_id, status, env, time_taken)
+            request_logs.complete_request(request_id, status, env, time_taken, query_count)
             expect(logger).not_to have_received(:info)
           end
         end

spec/unit/lib/sequel/extensions/request_query_counter_spec.rb

@@ -0,0 +1,15 @@
+require 'spec_helper'
+
+RSpec.describe Sequel::RequestQueryCounter do
+  let(:db) { DbConfig.new.connection }
+
+  before do
+    db.extension(:request_query_counter)
+  end
+
+  it 'increments the thread local query count on each query' do
+    expect do
+      db.fetch('SELECT 1').all
+    end.to change(VCAP::Request, :db_query_count).by(1)
+  end
+end

spec/unit/middleware/request_logs_spec.rb

@@ -21,7 +21,7 @@ module Middleware
 
         it 'calls complete request on request logs after the request' do
           middleware.call(env)
-          expect(request_logs).to have_received(:complete_request).with('ID', 200, { 'cf.request_id' => 'ID' }, be_a(Numeric))
+          expect(request_logs).to have_received(:complete_request).with('ID', 200, { 'cf.request_id' => 'ID' }, be_a(Numeric), be_a(Numeric))
         end
       end
     end

spec/unit/middleware/request_metrics_spec.rb

@@ -9,7 +9,8 @@ module Middleware
       let(:request_metrics) { instance_double(VCAP::CloudController::Metrics::RequestMetrics, start_request: nil, complete_request: nil) }
 
       describe 'handling the request' do
-        it 'calls start request on request metrics before the request' do
+        it 'resets the db query count and calls start request on request metrics before the request' do
+          expect(VCAP::Request).to receive(:reset_db_query_count)
           middleware.call({})
           expect(request_metrics).to have_received(:start_request)
         end