122 lines
		
	
	
		
			3.6 KiB
		
	
	
	
		
			Ruby
		
	
	
	
			
		
		
	
	
			122 lines
		
	
	
		
			3.6 KiB
		
	
	
	
		
			Ruby
		
	
	
	
| # 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::*`
 | |
|     #
 | |
|     # Sometimes this might cause errors in specs.
 | |
|     # This is best to be disable with `describe '...', query_analyzers: false do`
 | |
|     class QueryAnalyzer
 | |
|       include ::Singleton
 | |
| 
 | |
|       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
 | |
|           # to avoid stack too deep this detects recursive call of subscriber
 | |
|           with_ignored_recursive_calls do
 | |
|             process_sql(event.payload[:sql], event.payload[:connection])
 | |
|           end
 | |
|         end
 | |
|       end
 | |
| 
 | |
|       def within(analyzers = all_analyzers)
 | |
|         newly_enabled_analyzers = begin!(analyzers)
 | |
| 
 | |
|         begin
 | |
|           yield
 | |
|         ensure
 | |
|           end!(newly_enabled_analyzers)
 | |
|         end
 | |
|       end
 | |
| 
 | |
|       # Enable query analyzers (only the ones that were not yet enabled)
 | |
|       # Returns a list of newly enabled analyzers
 | |
|       def begin!(analyzers)
 | |
|         analyzers.select do |analyzer|
 | |
|           next if enabled_analyzers.include?(analyzer)
 | |
| 
 | |
|           if analyzer.enabled?
 | |
|             analyzer.begin!
 | |
|             enabled_analyzers.append(analyzer)
 | |
| 
 | |
|             true
 | |
|           end
 | |
|         rescue StandardError, ::Gitlab::Database::QueryAnalyzers::Base::QueryAnalyzerError => e
 | |
|           Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e)
 | |
| 
 | |
|           false
 | |
|         end
 | |
|       end
 | |
| 
 | |
|       # Disable enabled query analyzers (only the ones that were enabled previously)
 | |
|       def end!(analyzers)
 | |
|         analyzers.each do |analyzer|
 | |
|           next unless enabled_analyzers.delete(analyzer)
 | |
| 
 | |
|           analyzer.end!
 | |
|         rescue StandardError, ::Gitlab::Database::QueryAnalyzers::Base::QueryAnalyzerError => e
 | |
|           Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e)
 | |
|         end
 | |
|       end
 | |
| 
 | |
|       private
 | |
| 
 | |
|       def enabled_analyzers
 | |
|         Thread.current[:query_analyzer_enabled_analyzers] ||= []
 | |
|       end
 | |
| 
 | |
|       def process_sql(sql, connection)
 | |
|         analyzers = enabled_analyzers
 | |
|         return unless analyzers&.any?
 | |
| 
 | |
|         parsed = parse(sql, connection)
 | |
|         return unless parsed
 | |
| 
 | |
|         analyzers.each do |analyzer|
 | |
|           next if analyzer.suppressed? && !analyzer.requires_tracking?(parsed)
 | |
| 
 | |
|           analyzer.analyze(parsed)
 | |
|         rescue StandardError, ::Gitlab::Database::QueryAnalyzers::Base::QueryAnalyzerError => 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 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
 | |
| 
 | |
|       def with_ignored_recursive_calls
 | |
|         return if Thread.current[:query_analyzer_recursive]
 | |
| 
 | |
|         begin
 | |
|           Thread.current[:query_analyzer_recursive] = true
 | |
|           yield
 | |
|         ensure
 | |
|           Thread.current[:query_analyzer_recursive] = nil
 | |
|         end
 | |
|       end
 | |
|     end
 | |
|   end
 | |
| end
 |