1 files changed, 66 insertions, 0 deletions
diff --git a/lib/gitlab/database/transaction/observer.rb b/lib/gitlab/database/transaction/observer.rb
new file mode 100644
index 00000000000..7888f0916e3
--- /dev/null
+++ b/lib/gitlab/database/transaction/observer.rb
@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Transaction
+      class Observer
+        INSTRUMENTED_STATEMENTS = %w[BEGIN SAVEPOINT ROLLBACK RELEASE].freeze
+        LONGEST_COMMAND_LENGTH = 'ROLLBACK TO SAVEPOINT'.length
+        START_COMMENT = '/*'
+        END_COMMENT = '*/'
+
+        def self.instrument_transactions(cmd, event)
+          connection = event.payload[:connection]
+          manager = connection&.transaction_manager
+          return unless manager.respond_to?(:transaction_context)
+
+          context = manager.transaction_context
+          return if context.nil?
+
+          if cmd.start_with?('BEGIN')
+            context.set_start_time
+            context.set_depth(0)
+            context.track_sql(event.payload[:sql])
+          elsif cmd.start_with?('SAVEPOINT ')
+            context.set_depth(manager.open_transactions)
+            context.increment_savepoints
+          elsif cmd.start_with?('ROLLBACK TO SAVEPOINT')
+            context.increment_rollbacks
+          elsif cmd.start_with?('RELEASE SAVEPOINT ')
+            context.increment_releases
+          end
+        end
+
+        def self.register!
+          ActiveSupport::Notifications.subscribe('sql.active_record') do |event|
+            sql = event.payload.dig(:sql).to_s
+            cmd = extract_sql_command(sql)
+
+            if cmd.start_with?(*INSTRUMENTED_STATEMENTS)
+              self.instrument_transactions(cmd, event)
+            end
+          end
+        end
+
+        def self.extract_sql_command(sql)
+          return sql unless sql.start_with?(START_COMMENT)
+
+          index = sql.index(END_COMMENT)
+
+          return sql unless index
+
+          # /* comment */ SELECT
+          #
+          # We offset using a position of the end comment + 1 character to
+          # accomodate a space between Marginalia comment and a SQL statement.
+          offset = index + END_COMMENT.length + 1
+
+          # Avoid duplicating the entire string. This isn't optimized to
+          # strip extra spaces, but we assume that this doesn't happen
+          # for performance reasons.
+          sql[offset..offset + LONGEST_COMMAND_LENGTH]
+        end
+      end
+    end
+  end
+end