Add query count metrics

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

@@ -9,13 +9,15 @@ def initialize(app, request_logs)
       def call(env)
         request_id = env['cf.request_id']
         @request_logs.start_request(request_id, env)
+        VCAP::Request.reset_db_query_count
         start_time = Time.now
 
         status, headers, body = @app.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

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

@@ -46,10 +46,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, nil)
             expect(logger).to have_received(:info).with(
               /\ACompleted 200 vcap-request-id: ID/,
               time_taken_in_ms: 30,
+              db_query_count: nil,
               request_method: 'request_method',
               request_fullpath: 'filtered_path',
               status_code: 200
@@ -59,7 +60,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, nil)
             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
+    VCAP::Request.reset_db_query_count
+    db.fetch('SELECT 1').all
+    expect(VCAP::Request.db_query_count).to eq(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), 0)
         end
       end
     end