1 files changed, 36 insertions, 20 deletions
diff --git a/lib/gitlab/metrics/subscribers/active_record.rb b/lib/gitlab/metrics/subscribers/active_record.rb
index d725d8d7b29..5eefef02507 100644
--- a/lib/gitlab/metrics/subscribers/active_record.rb
+++ b/lib/gitlab/metrics/subscribers/active_record.rb
@@ -9,6 +9,16 @@ module Gitlab
 
         IGNORABLE_SQL = %w{BEGIN COMMIT}.freeze
         DB_COUNTERS = %i{db_count db_write_count db_cached_count}.freeze
+        SQL_COMMANDS_WITH_COMMENTS_REGEX = /\A(\/\*.*\*\/\s)?((?!(.*[^\w'"](DELETE|UPDATE|INSERT INTO)[^\w'"])))(WITH.*)?(SELECT)((?!(FOR UPDATE|FOR SHARE)).)*$/i.freeze
+
+        DURATION_BUCKET = [0.05, 0.1, 0.25].freeze
+
+        # This event is published from ActiveRecordBaseTransactionMetrics and
+        # used to record a database transaction duration when calling
+        # ActiveRecord::Base.transaction {} block.
+        def transaction(event)
+          observe(:gitlab_database_transaction_seconds, event)
+        end
 
         def sql(event)
           # Mark this thread as requiring a database connection. This is used
@@ -17,51 +27,57 @@ module Gitlab
           Thread.current[:uses_db_connection] = true
 
           payload = event.payload
-          return if payload[:name] == 'SCHEMA' || IGNORABLE_SQL.include?(payload[:sql])
+          return if ignored_query?(payload)
 
-          increment_db_counters(payload)
+          increment(:db_count)
+          increment(:db_cached_count) if cached_query?(payload)
+          increment(:db_write_count) unless select_sql_command?(payload)
 
-          current_transaction&.observe(:gitlab_sql_duration_seconds, event.duration / 1000.0) do
-            buckets [0.05, 0.1, 0.25]
-          end
+          observe(:gitlab_sql_duration_seconds, event)
         end
 
         def self.db_counter_payload
           return {} unless Gitlab::SafeRequestStore.active?
 
-          DB_COUNTERS.map do |counter|
-            [counter, Gitlab::SafeRequestStore[counter].to_i]
-          end.to_h
+          payload = {}
+          DB_COUNTERS.each do |counter|
+            payload[counter] = Gitlab::SafeRequestStore[counter].to_i
+          end
+          payload
         end
 
         private
 
-        def select_sql_command?(payload)
-          payload[:sql].match(/\A((?!(.*[^\w'"](DELETE|UPDATE|INSERT INTO)[^\w'"])))(WITH.*)?(SELECT)((?!(FOR UPDATE|FOR SHARE)).)*$/i)
+        def ignored_query?(payload)
+          payload[:name] == 'SCHEMA' || IGNORABLE_SQL.include?(payload[:sql])
         end
 
-        def increment_db_counters(payload)
-          increment(:db_count)
-
-          if payload.fetch(:cached, payload[:name] == 'CACHE')
-            increment(:db_cached_count)
-          end
+        def cached_query?(payload)
+          payload.fetch(:cached, payload[:name] == 'CACHE')
+        end
 
-          increment(:db_write_count) unless select_sql_command?(payload)
+        def select_sql_command?(payload)
+          payload[:sql].match(SQL_COMMANDS_WITH_COMMENTS_REGEX)
         end
 
         def increment(counter)
           current_transaction&.increment("gitlab_transaction_#{counter}_total".to_sym, 1)
 
-          if Gitlab::SafeRequestStore.active?
-            Gitlab::SafeRequestStore[counter] = Gitlab::SafeRequestStore[counter].to_i + 1
+          Gitlab::SafeRequestStore[counter] = Gitlab::SafeRequestStore[counter].to_i + 1
+        end
+
+        def observe(histogram, event)
+          current_transaction&.observe(histogram, event.duration / 1000.0) do
+            buckets DURATION_BUCKET
           end
         end
 
         def current_transaction
-          ::Gitlab::Metrics::Transaction.current
+          ::Gitlab::Metrics::WebTransaction.current || ::Gitlab::Metrics::BackgroundTransaction.current
         end
       end
     end
   end
 end
+
+Gitlab::Metrics::Subscribers::ActiveRecord.prepend_if_ee('EE::Gitlab::Metrics::Subscribers::ActiveRecord')