Observer to add query statistics

Relates to
https://gitlab.com/gitlab-org/database-team/team-tasks/-/issues/125

lib/gitlab/database/migrations/observation.rb

@@ -7,7 +7,8 @@ module Gitlab
         :migration,
         :walltime,
         :success,
-        :total_database_size_change
+        :total_database_size_change,
+        :query_statistics
       )
     end
   end

lib/gitlab/database/migrations/observers.rb

@@ -6,7 +6,8 @@ module Gitlab
       module Observers
         def self.all_observers
           [
-            TotalDatabaseSizeChange.new
+            TotalDatabaseSizeChange.new,
+            QueryStatistics.new
           ]
         end
       end

lib/gitlab/database/migrations/observers/query_statistics.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Migrations
+      module Observers
+        # This observer gathers statistics from the pg_stat_statements extension.
+        # Notice that this extension is not installed by default. In case it cannot
+        # be found, the observer does nothing and doesn't throw an error.
+        class QueryStatistics < MigrationObserver
+          include Gitlab::Database::SchemaHelpers
+
+          def before
+            return unless enabled?
+
+            connection.execute('select pg_stat_statements_reset()')
+          end
+
+          def record(observation)
+            return unless enabled?
+
+            observation.query_statistics = connection.execute(<<~SQL)
+              SELECT query, calls, total_time, max_time, mean_time, rows
+              FROM pg_stat_statements
+              ORDER BY total_time DESC
+            SQL
+          end
+
+          private
+
+          def enabled?
+            function_exists?(:pg_stat_statements_reset) && connection.view_exists?(:pg_stat_statements)
+          end
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/database/migrations/observers/query_statistics_spec.rb

+# frozen_string_literal: true
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Migrations::Observers::QueryStatistics do
+  subject { described_class.new }
+
+  let(:connection) { ActiveRecord::Base.connection }
+
+  def mock_pgss(enabled: true)
+    if enabled
+      allow(subject).to receive(:function_exists?).with(:pg_stat_statements_reset).and_return(true)
+      allow(connection).to receive(:view_exists?).with(:pg_stat_statements).and_return(true)
+    else
+      allow(subject).to receive(:function_exists?).with(:pg_stat_statements_reset).and_return(false)
+      allow(connection).to receive(:view_exists?).with(:pg_stat_statements).and_return(false)
+    end
+  end
+
+  describe '#before' do
+    context 'with pgss available' do
+      it 'resets pg_stat_statements' do
+        mock_pgss(enabled: true)
+        expect(connection).to receive(:execute).with('select pg_stat_statements_reset()').once
+
+        subject.before
+      end
+    end
+
+    context 'without pgss available' do
+      it 'executes nothing' do
+        mock_pgss(enabled: false)
+        expect(connection).not_to receive(:execute)
+
+        subject.before
+      end
+    end
+  end
+
+  describe '#record' do
+    let(:observation) { Gitlab::Database::Migrations::Observation.new }
+    let(:result) { double }
+    let(:pgss_query) do
+      <<~SQL
+        SELECT query, calls, total_time, max_time, mean_time, rows
+        FROM pg_stat_statements
+        ORDER BY total_time DESC
+      SQL
+    end
+
+    context 'with pgss available' do
+      it 'fetches data from pg_stat_statements and stores on the observation' do
+        mock_pgss(enabled: true)
+        expect(connection).to receive(:execute).with(pgss_query).once.and_return(result)
+
+        expect { subject.record(observation) }.to change { observation.query_statistics }.from(nil).to(result)
+      end
+    end
+
+    context 'without pgss available' do
+      it 'executes nothing' do
+        mock_pgss(enabled: false)
+        expect(connection).not_to receive(:execute)
+
+        expect { subject.record(observation) }.not_to change { observation.query_statistics }
+      end
+    end
+  end
+end