Add `QueryAnalyzer` to have a single approach to hook all analyzers

This adds a generic and central place to parse and normalize queries
via `pg_query`. Makes it easy to write new analyzers and hook them depending
on other conditions.

config/initializers/database_query_analyzers.rb

+# frozen_string_literal: true
+
+# Currently we register validator only for `dev` or `test` environment
+Gitlab::Database::QueryAnalyzer.new.hook! if Gitlab.dev_or_test_env?

lib/gitlab/database/load_balancing/configuration.rb

@@ -77,6 +77,10 @@ module Gitlab
           (@primary_model || @model).connection_specification_name
         end
 
+        def primary_db_config
+          (@primary_model || @model).connection_db_config
+        end
+
         def replica_db_config
           @model.connection_db_config
         end

lib/gitlab/database/query_analyzer.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    # The purpose of this class is to implement a various query analyzers based on `pg_query`
+    # And process them all via `Gitlab::Database::QueryAnalyzers::*`
+    class QueryAnalyzer
+      ANALYZERS = [].freeze
+
+      Parsed = Struct.new(
+        :sql, :connection, :pg
+      )
+
+      def hook!
+        @subscriber = ActiveSupport::Notifications.subscribe('sql.active_record') do |event|
+          process_sql(event.payload[:sql], event.payload[:connection])
+        end
+      end
+
+      private
+
+      def process_sql(sql, connection)
+        analyzers = enabled_analyzers(connection)
+        return unless analyzers.any?
+
+        parsed = parse(sql, connection)
+        return unless parsed
+
+        analyzers.each do |analyzer|
+          analyzer.analyze(parsed)
+        rescue => e # rubocop:disable Style/RescueStandardError
+          # We catch all standard errors to prevent validation errors to introduce fatal errors in production
+          Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e)
+        end
+      end
+
+      def enabled_analyzers(connection)
+        ANALYZERS.select do |analyzer|
+          analyzer.enabled?(connection)
+        rescue StandardError => e # rubocop:disable Style/RescueStandardError
+          # We catch all standard errors to prevent validation errors to introduce fatal errors in production
+          Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e)
+        end
+      end
+
+      def parse(sql, connection)
+        parsed = PgQuery.parse(sql)
+        return unless parsed
+
+        normalized = PgQuery.normalize(sql)
+        Parsed.new(normalized, connection, parsed)
+      rescue PgQuery::ParseError => e
+        # Ignore PgQuery parse errors (due to depth limit or other reasons)
+        Gitlab::ErrorTracking.track_exception(e)
+
+        nil
+      end
+    end
+  end
+end

lib/gitlab/database/query_analyzers/base.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module QueryAnalyzers
+      class Base
+        def self.enabled?(connection)
+          raise NotImplementedError
+        end
+
+        def self.analyze(parsed)
+          raise NotImplementedError
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/database/query_analyzer_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::QueryAnalyzer do
+  let(:analyzer) { double(:query_analyzer) }
+
+  before do
+    stub_const('Gitlab::Database::QueryAnalyzer::ANALYZERS', [analyzer])
+  end
+
+  context 'the hook is enabled by default in specs' do
+    it 'does process queries and gets normalized SQL' do
+      expect(analyzer).to receive(:enabled?).and_return(true)
+      expect(analyzer).to receive(:analyze) do |parsed|
+        expect(parsed.sql).to include("SELECT $1 FROM projects")
+        expect(parsed.pg.tables).to eq(%w[projects])
+      end
+
+      Project.connection.execute("SELECT 1 FROM projects")
+    end
+  end
+
+  describe '#process_sql' do
+    it 'does not analyze query if not enabled' do
+      expect(analyzer).to receive(:enabled?).and_return(false)
+      expect(analyzer).not_to receive(:analyze)
+
+      process_sql("SELECT 1 FROM projects")
+    end
+
+    it 'does analyze query if enabled' do
+      expect(analyzer).to receive(:enabled?).and_return(true)
+      expect(analyzer).to receive(:analyze) do |parsed|
+        expect(parsed.sql).to eq("SELECT $1 FROM projects")
+        expect(parsed.pg.tables).to eq(%w[projects])
+      end
+
+      process_sql("SELECT 1 FROM projects")
+    end
+
+    it 'does track exception if query cannot be parsed' do
+      expect(analyzer).to receive(:enabled?).and_return(true)
+      expect(analyzer).not_to receive(:analyze)
+      expect(Gitlab::ErrorTracking).to receive(:track_exception)
+
+      expect { process_sql("invalid query") }.not_to raise_error
+    end
+
+    it 'does track exception if analyzer raises exception on enabled?' do
+      expect(analyzer).to receive(:enabled?).and_raise('exception')
+      expect(analyzer).not_to receive(:analyze)
+      expect(Gitlab::ErrorTracking).to receive(:track_and_raise_for_dev_exception)
+
+      expect { process_sql("SELECT 1 FROM projects") }.not_to raise_error
+    end
+
+    it 'does track exception if analyzer raises exception on analyze' do
+      expect(analyzer).to receive(:enabled?).and_return(true)
+      expect(analyzer).to receive(:analyze).and_raise('exception')
+      expect(Gitlab::ErrorTracking).to receive(:track_and_raise_for_dev_exception)
+
+      expect { process_sql("SELECT 1 FROM projects") }.not_to raise_error
+    end
+
+    def process_sql(sql)
+      ApplicationRecord.connection.load_balancer.read_write do |connection|
+        described_class.new.send(:process_sql, sql, connection)
+      end
+    end
+  end
+end