Store Sherlock::Query in Peek adapter

Signed-off-by: Rémy Coutable <remy@rymai.me>

3 files changed, 104 insertions, 0 deletions

diff --git a/config/initializers/peek.rb b/config/initializers/peek.rb
index 73da7be7889..6cb0e711e48 100644
--- a/config/initializers/peek.rb
+++ b/config/initializers/peek.rb
@@ -7,3 +7,35 @@ Peek.into Peek::Views::Redis
 Peek.into Peek::Views::Sidekiq
 Peek.into Peek::Views::Rblineprof
 Peek.into Peek::Views::GC
+
+if Gitlab::Database.mysql?
+  class Mysql2::Client
+    class << self
+      attr_accessor :query_details
+    end
+    self.query_details = Concurrent::Array.new
+  end
+
+  module Peek
+    module Views
+      class Mysql2 < View
+        prepend ::Gitlab::PerformanceBar::PeekMysqlWithQueries
+      end
+    end
+  end
+else
+  class PG::Connection
+    class << self
+      attr_accessor :query_details
+    end
+    self.query_details = Concurrent::Array.new
+  end
+
+  module Peek
+    module Views
+      class PG < View
+        prepend ::Gitlab::PerformanceBar::PeekPgWithQueries
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/performance_bar/peek_mysql_with_queries.rb b/lib/gitlab/performance_bar/peek_mysql_with_queries.rb
new file mode 100644
index 00000000000..ea577de5a69
--- /dev/null
+++ b/lib/gitlab/performance_bar/peek_mysql_with_queries.rb
@@ -0,0 +1,36 @@
+# Inspired by https://github.com/peek/peek-mysql2/blob/master/lib/peek/views/mysql2.rb
+module Gitlab
+  module PerformanceBar
+    module PeekMysqlWithQueries
+      def queries
+        ::Mysql2::Client.query_details
+      end
+
+      def results
+        super.merge(queries: queries)
+      end
+
+      private
+
+      def setup_subscribers
+        super
+
+        # Reset each counter when a new request starts
+        before_request do
+          ::Mysql2::Client.query_details = []
+        end
+
+        subscribe('sql.active_record') do |_, start, finish, _, data|
+          if RequestStore.active? && RequestStore.store[:peek_enabled]
+            track_query(data[:sql].strip, data[:binds], start, finish)
+          end
+        end
+      end
+
+      def track_query(raw_query, bindings, start, finish)
+        query = Gitlab::Sherlock::Query.new(raw_query, start, finish)
+        ::Mysql2::Client.query_details << query.formatted_query
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/performance_bar/peek_pg_with_queries.rb b/lib/gitlab/performance_bar/peek_pg_with_queries.rb
new file mode 100644
index 00000000000..6e9cb593d36
--- /dev/null
+++ b/lib/gitlab/performance_bar/peek_pg_with_queries.rb
@@ -0,0 +1,36 @@
+# Inspired by https://github.com/peek/peek-pg/blob/master/lib/peek/views/pg.rb
+module Gitlab
+  module PerformanceBar
+    module PeekPgWithQueries
+      def queries
+        ::PG::Connection.query_details
+      end
+
+      def results
+        super.merge(queries: queries)
+      end
+
+      private
+
+      def setup_subscribers
+        super
+
+        # Reset each counter when a new request starts
+        before_request do
+          ::PG::Connection.query_details = []
+        end
+
+        subscribe('sql.active_record') do |_, start, finish, _, data|
+          if RequestStore.active? && RequestStore.store[:peek_enabled]
+            track_query(data[:sql].strip, data[:binds], start, finish)
+          end
+        end
+      end
+
+      def track_query(raw_query, bindings, start, finish)
+        query = Gitlab::Sherlock::Query.new(raw_query, start, finish)
+        ::PG::Connection.query_details << query.formatted_query
+      end
+    end
+  end
+end