Implement review feedback

config/feature_flags/development/ci_pipeline_logger_sql.yml → config/feature_flags/development/ci_pipeline_logger_sql_count.yml

 ---
-name: ci_pipeline_logger_sql
+name: ci_pipeline_logger_sql_count
 introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/77117
 rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/346246
 milestone: '14.7'

lib/gitlab/ci/pipeline/chain/create.rb

@@ -9,7 +9,7 @@ module Gitlab
           include Gitlab::Utils::StrongMemoize
 
           def perform!
-            logger.instrument(:pipeline_save, sql: true) do
+            logger.instrument_with_sql(:pipeline_save) do
               BulkInsertableAssociations.with_bulk_insert do
                 tags = extract_tag_list_by_status
 

lib/gitlab/ci/pipeline/logger.rb

@@ -23,18 +23,28 @@ module Gitlab
           log_conditions.push(block)
         end
 
-        def instrument(operation, sql: false)
+        def instrument(operation)
           return yield unless enabled?
 
           raise ArgumentError, 'block not given' unless block_given?
 
-          op_start_db_counters = current_db_counter_payload if sql && sql_logging_enabled?
           op_started_at = current_monotonic_time
 
           result = yield
 
           observe("#{operation}_duration_s", current_monotonic_time - op_started_at)
-          observe_sql(operation, op_start_db_counters) if sql && sql_logging_enabled?
+
+          result
+        end
+
+        def instrument_with_sql(operation)
+          return instrument(operation) { yield } unless sql_logging_enabled?
+
+          op_start_db_counters = current_db_counter_payload
+
+          result = instrument(operation) { yield }
+
+          observe_sql_counters(operation, op_start_db_counters, current_db_counter_payload)
 
           result
         end
@@ -100,8 +110,8 @@ module Gitlab
           @observations ||= Hash.new { |hash, key| hash[key] = [] }
         end
 
-        def observe_sql(operation, start_db_counters)
-          current_db_counter_payload.each do |key, value|
+        def observe_sql_counters(operation, start_db_counters, end_db_counters)
+          end_db_counters.each do |key, value|
             result = value - start_db_counters.fetch(key, 0)
             next if result == 0
 
@@ -115,7 +125,7 @@ module Gitlab
 
         def sql_logging_enabled?
           strong_memoize(:sql_logging_enabled) do
-            ::Feature.enabled?(:ci_pipeline_logger_sql, project, default_enabled: :yaml)
+            ::Feature.enabled?(:ci_pipeline_logger_sql_count, project, default_enabled: :yaml)
           end
         end
       end

spec/lib/gitlab/ci/pipeline/logger_spec.rb

@@ -39,36 +39,88 @@ RSpec.describe ::Gitlab::Ci::Pipeline::Logger do
       expect { logger.instrument(:expensive_operation) }
         .to raise_error(ArgumentError, 'block not given')
     end
+  end
+
+  describe '#instrument_with_sql', :request_store do
+    subject(:instrument_with_sql) do
+      logger.instrument_with_sql(:expensive_operation, &operation)
+    end
+
+    def loggable_data(count:, db_count: nil)
+      keys = %w[
+        expensive_operation_duration_s
+        expensive_operation_db_count
+        expensive_operation_db_primary_count
+        expensive_operation_db_primary_duration_s
+        expensive_operation_db_main_count
+        expensive_operation_db_main_duration_s
+      ]
+
+      data = keys.each.with_object({}) do |key, accumulator|
+        accumulator[key] = {
+          'count' => count,
+          'avg' => a_kind_of(Numeric),
+          'max' => a_kind_of(Numeric),
+          'min' => a_kind_of(Numeric)
+        }
+      end
 
-    context 'with SQL instrumentation', :request_store do
-      it "returns the block's value" do
-        expect(logger.instrument(:expensive_operation, sql: true) { Project.count }).to eq(Project.count)
+      if db_count
+        data['expensive_operation_db_count']['max'] = db_count
+        data['expensive_operation_db_count']['min'] = db_count
+        data['expensive_operation_db_count']['avg'] = db_count
       end
 
+      data
+    end
+
+    context 'with a single query' do
+      let(:operation) { -> { Project.count } }
+
+      it { is_expected.to eq(operation.call) }
+
       it 'includes SQL metrics' do
-        loggable_data = %w[
-          expensive_operation_duration_s
-          expensive_operation_db_count
-          expensive_operation_db_primary_count
-          expensive_operation_db_primary_duration_s
-          expensive_operation_db_main_count
-          expensive_operation_db_main_duration_s
-        ].each.with_object({}) do |key, accumulator|
-          accumulator[key] = {
-            'count' => 1,
-            'avg' => a_kind_of(Numeric),
-            'max' => a_kind_of(Numeric),
-            'min' => a_kind_of(Numeric)
-          }
-        end
+        instrument_with_sql
+
+        expect(logger.observations_hash)
+          .to match(a_hash_including(loggable_data(count: 1, db_count: 1)))
+      end
+    end
+
+    context 'with multiple queries' do
+      let(:operation) { -> { Ci::Build.count + Ci::Bridge.count } }
+
+      it { is_expected.to eq(operation.call) }
+
+      it 'includes SQL metrics' do
+        instrument_with_sql
+
+        expect(logger.observations_hash)
+          .to match(a_hash_including(loggable_data(count: 1, db_count: 2)))
+      end
+    end
+
+    context 'with multiple observations' do
+      let(:operation) { -> { Ci::Build.count + Ci::Bridge.count } }
+
+      it 'includes SQL metrics' do
+        2.times { logger.instrument_with_sql(:expensive_operation, &operation) }
 
-        logger.instrument(:expensive_operation, sql: true) { Project.count }
-        expect(logger.observations_hash).to match(a_hash_including(loggable_data))
+        expect(logger.observations_hash)
+          .to match(a_hash_including(loggable_data(count: 2, db_count: 2)))
       end
+    end
+
+    context 'when there are not SQL operations' do
+      let(:operation) { -> { 123 } }
+
+      it { is_expected.to eq(operation.call) }
+
+      it 'does not include SQL metrics' do
+        instrument_with_sql
 
-      it 'does not include SQL metrics when there is no SQL operation' do
-        logger.instrument(:expensive_operation, sql: true) { 123 }
-        expect(logger.observations_hash.keys).to match_array(['expensive_operation_duration_s'])
+        expect(logger.observations_hash.keys)
+          .to match_array(['expensive_operation_duration_s'])
       end
     end
   end