Make QueryAnalyzers to hold a state and hook to middleware

One of the problems of QueryAnalyzers is that their state
is evaluated each time. This is problematic for various reasons
as we need to re-evaluate feature flags, or other dynamic conditions.

This makes QueryAnalyzers behavior sticky. The enabled flag is checked
exactly once for each request/worker/spec. If it resolves to true,
the query analyzer will be enabled for the whole execution of a job.

This makes also QueryAnalyzer to hold a context that is set/reset
on boundaries of job execution.

config/initializers/database_query_analyzers.rb

@@ -3,4 +3,9 @@
 # Currently we register validator only for `dev` or `test` environment
 if Gitlab.dev_or_test_env? || Gitlab::Utils.to_boolean('GITLAB_ENABLE_QUERY_ANALYZERS', default: false)
   Gitlab::Database::QueryAnalyzer.instance.hook!
+
+  Gitlab::Application.configure do |config|
+    # ApolloUploadServer::Middleware expects to find uploaded files ready to use
+    config.middleware.use(Gitlab::Middleware::QueryAnalyzer)
+  end
 end

lib/gitlab/database/query_analyzer.rb

@@ -4,15 +4,22 @@ 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::*`
+    #
+    # Sometimes this might cause errors in specs.
+    # This is best to be disable with `describe '...', query_analyzers: false do`
     class QueryAnalyzer
       include ::Singleton
 
-      ANALYZERS = [].freeze
-
       Parsed = Struct.new(
         :sql, :connection, :pg
       )
 
+      attr_reader :all_analyzers
+
+      def initialize
+        @all_analyzers = []
+      end
+
       def hook!
         @subscriber = ActiveSupport::Notifications.subscribe('sql.active_record') do |event|
           # In some cases analyzer code might trigger another SQL call
@@ -23,30 +30,63 @@ module Gitlab
         end
       end
 
-      private
+      def within
+        return yield if enabled_analyzers
+
+        begin!
+
+        begin
+          yield
+        ensure
+          end!
+        end
+      end
 
       def process_sql(sql, connection)
-        analyzers = enabled_analyzers(connection)
-        return unless analyzers.any?
+        analyzers = enabled_analyzers
+        return unless analyzers&.any?
 
         parsed = parse(sql, connection)
         return unless parsed
 
         analyzers.each do |analyzer|
           analyzer.analyze(parsed)
-        rescue => e # rubocop:disable Style/RescueStandardError
+        rescue StandardError => e
           # 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
+      private
+
+      # Enable query analyzers
+      def begin!
+        analyzers = all_analyzers.select do |analyzer|
+          if analyzer.enabled?
+            analyzer.begin!
+
+            true
+          end
+        rescue StandardError => e
+          Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e)
+        end
+
+        Thread.current[:query_analyzer_enabled_analyzers] = analyzers
+      end
+
+      # Disable enabled query analyzers
+      def end!
+        enabled_analyzers.select do |analyzer|
+          analyzer.end!
+        rescue StandardError => e
           Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e)
         end
+
+        Thread.current[:query_analyzer_enabled_analyzers] = nil
+      end
+
+      def enabled_analyzers
+        Thread.current[:query_analyzer_enabled_analyzers]
       end
 
       def parse(sql, connection)

lib/gitlab/database/query_analyzers/base.rb

@@ -4,7 +4,19 @@ module Gitlab
   module Database
     module QueryAnalyzers
       class Base
-        def self.enabled?(connection)
+        def self.begin!
+          Thread.current[self.class.name] = {}
+        end
+
+        def self.end!
+          Thread.current[self.class.name] = nil
+        end
+
+        def self.context
+          Thread.current[self.class.name]
+        end
+
+        def self.enabled?
           raise NotImplementedError
         end
 

lib/gitlab/middleware/query_analyzer.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Middleware
+    class QueryAnalyzer
+      def initialize(app)
+        @app = app
+      end
+
+      def call(env)
+        ::Gitlab::Database::QueryAnalyzer.instance.within { @app.call(env) }
+      end
+    end
+  end
+end

lib/gitlab/sidekiq_middleware.rb

@@ -33,6 +33,7 @@ module Gitlab
         chain.add ::Gitlab::SidekiqMiddleware::BatchLoader
         chain.add ::Gitlab::SidekiqMiddleware::InstrumentationLogger
         chain.add ::Gitlab::SidekiqMiddleware::AdminMode::Server
+        chain.add ::Gitlab::SidekiqMiddleware::QueryAnalyzer
         chain.add ::Gitlab::SidekiqVersioning::Middleware
         chain.add ::Gitlab::SidekiqStatus::ServerMiddleware
         chain.add ::Gitlab::SidekiqMiddleware::WorkerContext::Server

lib/gitlab/sidekiq_middleware/query_analyzer.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module SidekiqMiddleware
+    class QueryAnalyzer
+      def call(worker, job, queue)
+        ::Gitlab::Database::QueryAnalyzer.instance.within { yield }
+      end
+    end
+  end
+end

spec/lib/gitlab/database/query_analyzer_spec.rb

@@ -2,11 +2,16 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::Database::QueryAnalyzer do
+RSpec.describe Gitlab::Database::QueryAnalyzer, query_analyzers: false do
   let(:analyzer) { double(:query_analyzer) }
+  let(:disabled_analyzer) { double(:disabled_query_analyzer) }
 
   before do
-    stub_const('Gitlab::Database::QueryAnalyzer::ANALYZERS', [analyzer])
+    allow(described_class.instance).to receive(:all_analyzers).and_return([analyzer, disabled_analyzer])
+    allow(analyzer).to receive(:enabled?).and_return(true)
+    allow(analyzer).to receive(:begin!)
+    allow(analyzer).to receive(:end!)
+    allow(disabled_analyzer).to receive(:enabled?).and_return(false)
   end
 
   context 'the hook is enabled by default in specs' do
@@ -17,7 +22,9 @@ RSpec.describe Gitlab::Database::QueryAnalyzer do
         expect(parsed.pg.tables).to eq(%w[projects])
       end
 
-      Project.connection.execute("SELECT 1 FROM projects")
+      described_class.instance.within do
+        Project.connection.execute("SELECT 1 FROM projects")
+      end
     end
 
     it 'does prevent recursive execution' do
@@ -26,7 +33,46 @@ RSpec.describe Gitlab::Database::QueryAnalyzer do
         Project.connection.execute("SELECT 1 FROM projects")
       end
 
-      Project.connection.execute("SELECT 1 FROM projects")
+      described_class.instance.within do
+        Project.connection.execute("SELECT 1 FROM projects")
+      end
+    end
+  end
+
+  describe '#within' do
+    context 'when it is already initialized' do
+      around do |example|
+        described_class.instance.within do
+          example.run
+        end
+      end
+
+      it 'does not evaluate enabled? again do yield block' do
+        expect(analyzer).not_to receive(:enabled?)
+
+        expect { |b| described_class.instance.within(&b) }.to yield_control
+      end
+    end
+
+    context 'when initializer is enabled' do
+      before do
+        expect(analyzer).to receive(:enabled?).and_return(true)
+      end
+
+      it 'calls begin! and end!' do
+        expect(analyzer).to receive(:begin!)
+        expect(analyzer).to receive(:end!)
+
+        expect { |b| described_class.instance.within(&b) }.to yield_control
+      end
+
+      it 'when begin! raises the end! is not called' do
+        expect(analyzer).to receive(:begin!).and_raise('exception')
+        expect(analyzer).not_to receive(:end!)
+        expect(Gitlab::ErrorTracking).to receive(:track_and_raise_for_dev_exception)
+
+        expect { |b| described_class.instance.within(&b) }.to yield_control
+      end
     end
   end
 
@@ -72,9 +118,16 @@ RSpec.describe Gitlab::Database::QueryAnalyzer do
       expect { process_sql("SELECT 1 FROM projects") }.not_to raise_error
     end
 
+    it 'does call analyze only on enabled initializers' do
+      expect(analyzer).to receive(:analyze)
+      expect(disabled_analyzer).not_to receive(:analyze)
+
+      expect { process_sql("SELECT 1 FROM projects") }.not_to raise_error
+    end
+
     def process_sql(sql)
       ApplicationRecord.load_balancer.read_write do |connection|
-        described_class.instance.send(:process_sql, sql, connection)
+        described_class.instance.process_sql(sql, connection)
       end
     end
   end

spec/support/database/query_analyzer.rb

+# frozen_string_literal: true
+
+# With the usage of `describe '...', query_analyzers: false`
+# can be disabled selectively
+
+RSpec.configure do |config|
+  config.around do |example|
+    if example.metadata.fetch(:query_analyzers, true)
+      ::Gitlab::Database::QueryAnalyzer.instance.within { example.run }
+    else
+      example.run
+    end
+  end
+end