diff --git a/lib/gitlab/database/query_analyzer.rb b/lib/gitlab/database/query_analyzer.rb
index 0c78dda734c734b7c2d10ddc355a11a07828b88a..6f64d04270fdbddb223dbb6c6cd222c75c1a5c77 100644
--- a/lib/gitlab/database/query_analyzer.rb
+++ b/lib/gitlab/database/query_analyzer.rb
@@ -30,52 +30,25 @@ def hook!
end
end
- def within(user_analyzers = nil)
- # Due to singleton nature of analyzers
- # only an outer invocation of the `.within`
- # is allowed to initialize them
- if already_within?
- raise 'Query analyzers are already defined, cannot re-define them.' if user_analyzers
-
- return yield
- end
-
- begin!(user_analyzers || all_analyzers)
+ def within(analyzers = all_analyzers)
+ newly_enabled_analyzers = begin!(analyzers)
begin
yield
ensure
- end!
+ end!(newly_enabled_analyzers)
end
end
- def already_within?
- # If analyzers are set they are already configured
- !enabled_analyzers.nil?
- 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)
- 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
-
- # Enable query analyzers
- def begin!(analyzers = all_analyzers)
- analyzers = analyzers.select do |analyzer|
if analyzer.enabled?
analyzer.begin!
+ enabled_analyzers.append(analyzer)
true
end
@@ -84,25 +57,40 @@ def begin!(analyzers = all_analyzers)
false
end
-
- Thread.current[:query_analyzer_enabled_analyzers] = analyzers
end
- # Disable enabled query analyzers
- def end!
- enabled_analyzers.select do |analyzer|
+ # 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
-
- Thread.current[:query_analyzer_enabled_analyzers] = nil
end
private
def enabled_analyzers
- Thread.current[:query_analyzer_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)
diff --git a/spec/lib/gitlab/database/query_analyzer_spec.rb b/spec/lib/gitlab/database/query_analyzer_spec.rb
index 3b4cbc79de2609909e265d1518d7db07b1d236c4..0b849063562a25e2cf9be95223122b16dee9c3e5 100644
--- a/spec/lib/gitlab/database/query_analyzer_spec.rb
+++ b/spec/lib/gitlab/database/query_analyzer_spec.rb
@@ -4,7 +4,7 @@
RSpec.describe Gitlab::Database::QueryAnalyzer, query_analyzers: false do
let(:analyzer) { double(:query_analyzer) }
- let(:user_analyzer) { double(:query_analyzer) }
+ let(:user_analyzer) { double(:user_query_analyzer) }
let(:disabled_analyzer) { double(:disabled_query_analyzer) }
before do
@@ -49,14 +49,36 @@
end
end
- it 'does not evaluate enabled? again do yield block' do
- expect(analyzer).not_to receive(:enabled?)
+ it 'does initialize analyzer only once' do
+ expect(analyzer).to receive(:enabled?).once
+ expect(analyzer).to receive(:begin!).once
+ expect(analyzer).to receive(:end!).once
expect { |b| described_class.instance.within(&b) }.to yield_control
end
- it 'raises exception when trying to re-define analyzers' do
- expect { |b| described_class.instance.within([user_analyzer], &b) }.to raise_error /Query analyzers are already defined, cannot re-define them/
+ it 'does initialize user analyzer when enabled' do
+ expect(user_analyzer).to receive(:enabled?).and_return(true)
+ expect(user_analyzer).to receive(:begin!)
+ expect(user_analyzer).to receive(:end!)
+
+ expect { |b| described_class.instance.within([user_analyzer], &b) }.to yield_control
+ end
+
+ it 'does initialize user analyzer only once' do
+ expect(user_analyzer).to receive(:enabled?).and_return(false, true)
+ expect(user_analyzer).to receive(:begin!).once
+ expect(user_analyzer).to receive(:end!).once
+
+ expect { |b| described_class.instance.within([user_analyzer, user_analyzer, user_analyzer], &b) }.to yield_control
+ end
+
+ it 'does not initializer user analyzer when disabled' do
+ expect(user_analyzer).to receive(:enabled?).and_return(false)
+ expect(user_analyzer).not_to receive(:begin!)
+ expect(user_analyzer).not_to receive(:end!)
+
+ expect { |b| described_class.instance.within([user_analyzer], &b) }.to yield_control
end
end
@@ -162,7 +184,7 @@
def process_sql(sql)
described_class.instance.within do
ApplicationRecord.load_balancer.read_write do |connection|
- described_class.instance.process_sql(sql, connection)
+ described_class.instance.send(:process_sql, sql, connection)
end
end
end
diff --git a/spec/lib/gitlab/database/query_analyzers/gitlab_schemas_metrics_spec.rb b/spec/lib/gitlab/database/query_analyzers/gitlab_schemas_metrics_spec.rb
index b8c1ecd908984bf9e3b9560863e4841b0642992a..0d687db0f961009469c499c3c97eb0aecbcdc3a6 100644
--- a/spec/lib/gitlab/database/query_analyzers/gitlab_schemas_metrics_spec.rb
+++ b/spec/lib/gitlab/database/query_analyzers/gitlab_schemas_metrics_spec.rb
@@ -140,7 +140,7 @@
def process_sql(model, sql)
Gitlab::Database::QueryAnalyzer.instance.within do
# Skip load balancer and retrieve connection assigned to model
- Gitlab::Database::QueryAnalyzer.instance.process_sql(sql, model.retrieve_connection)
+ Gitlab::Database::QueryAnalyzer.instance.send(:process_sql, sql, model.retrieve_connection)
end
end
end
diff --git a/spec/lib/gitlab/database/query_analyzers/restrict_allowed_schemas_spec.rb b/spec/lib/gitlab/database/query_analyzers/restrict_allowed_schemas_spec.rb
index a2c7916fa01e8c7ef08f9517087a42e1640c91ba..261bef58bb6e9ca00e548675b03baa0fa41bfa21 100644
--- a/spec/lib/gitlab/database/query_analyzers/restrict_allowed_schemas_spec.rb
+++ b/spec/lib/gitlab/database/query_analyzers/restrict_allowed_schemas_spec.rb
@@ -155,7 +155,7 @@ def process_sql(sql, model = ActiveRecord::Base)
yield if block_given?
# Skip load balancer and retrieve connection assigned to model
- Gitlab::Database::QueryAnalyzer.instance.process_sql(sql, model.retrieve_connection)
+ Gitlab::Database::QueryAnalyzer.instance.send(:process_sql, sql, model.retrieve_connection)
end
end
end
diff --git a/spec/support/database/query_analyzer.rb b/spec/support/database/query_analyzer.rb
index 6d6627d54b90312caea2cb3376f89e72fb0de279..aaa1b3516a327a88905736956966501130a98a0f 100644
--- a/spec/support/database/query_analyzer.rb
+++ b/spec/support/database/query_analyzer.rb
@@ -6,13 +6,17 @@
RSpec.configure do |config|
config.before do |example|
if example.metadata.fetch(:query_analyzers, true)
- ::Gitlab::Database::QueryAnalyzer.instance.begin!
+ ::Gitlab::Database::QueryAnalyzer.instance.begin!(
+ ::Gitlab::Database::QueryAnalyzer.instance.all_analyzers
+ )
end
end
config.after do |example|
if example.metadata.fetch(:query_analyzers, true)
- ::Gitlab::Database::QueryAnalyzer.instance.end!
+ ::Gitlab::Database::QueryAnalyzer.instance.end!(
+ ::Gitlab::Database::QueryAnalyzer.instance.all_analyzers
+ )
end
end
end