diff --git a/Gemfile b/Gemfile
index 0ac80a104f5d9ba925229be4c69c39ae979fa08d..c3150bea86273091bebd9b9aea1a3b6ca722692b 100644
--- a/Gemfile
+++ b/Gemfile
@@ -22,6 +22,7 @@ gem 'rugged', '~> 0.28'
gem 'grape-path-helpers', '~> 1.1'
gem 'faraday', '~> 0.12'
+gem 'marginalia', '~> 1.8.0'
# Authentication libraries
gem 'devise', '~> 4.6'
diff --git a/Gemfile.lock b/Gemfile.lock
index 55674a79cc2d342981b87969feccfd60c064d78b..9913d2d1df9b63e911e45cfbda5f1e0b36ddb65c 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -596,6 +596,9 @@ GEM
mail_room (0.9.1)
marcel (0.3.3)
mimemagic (~> 0.3.2)
+ marginalia (1.8.0)
+ actionpack (>= 2.3)
+ activerecord (>= 2.3)
memoist (0.16.0)
memoizable (0.4.2)
thread_safe (~> 0.3, >= 0.3.1)
@@ -1245,6 +1248,7 @@ DEPENDENCIES
lograge (~> 0.5)
loofah (~> 2.2)
mail_room (~> 0.9.1)
+ marginalia (~> 1.8.0)
memory_profiler (~> 0.9)
method_source (~> 0.8)
mimemagic (~> 0.3.2)
diff --git a/config/initializers/0_marginalia.rb b/config/initializers/0_marginalia.rb
new file mode 100644
index 0000000000000000000000000000000000000000..f88a90854e3d11871d46ea7f8a582a2e798de669
--- /dev/null
+++ b/config/initializers/0_marginalia.rb
@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+require 'marginalia'
+
+::Marginalia::Comment.extend(::Gitlab::Marginalia::Comment)
+
+# Patch to modify 'Marginalia::ActiveRecordInstrumentation.annotate_sql' method with feature check.
+# Orignal Marginalia::ActiveRecordInstrumentation is included to ActiveRecord::ConnectionAdapters::PostgreSQLAdapter in the Marginalia Railtie.
+# Refer: https://github.com/basecamp/marginalia/blob/v1.8.0/lib/marginalia/railtie.rb#L67
+ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(Gitlab::Marginalia::ActiveRecordInstrumentation)
+
+Marginalia::Comment.components = [:application, :controller, :action, :correlation_id, :jid, :job_class, :line]
+
+Gitlab::Marginalia.set_application_name
+
+Gitlab::Marginalia.enable_sidekiq_instrumentation
+
+Gitlab::Marginalia.set_feature_cache
diff --git a/doc/development/README.md b/doc/development/README.md
index 6aeaf31ed2909451e4c7dedc56c03528f5dfca49..3a972c4c5885db1b9aae294aad41538d1bf4ff59 100644
--- a/doc/development/README.md
+++ b/doc/development/README.md
@@ -103,6 +103,10 @@ description: 'Learn how to contribute to GitLab.'
- [Swapping tables](swapping_tables.md)
- [Deleting migrations](deleting_migrations.md)
+### Debugging
+
+- Tracing the source of an SQL query using query comments with [Marginalia](database_query_comments.md)
+
### Best practices
- [Merge Request checklist](database_merge_request_checklist.md)
diff --git a/doc/development/database_query_comments.md b/doc/development/database_query_comments.md
new file mode 100644
index 0000000000000000000000000000000000000000..c1a927ef234a338dfa03c0fd86d3c7cf093e570f
--- /dev/null
+++ b/doc/development/database_query_comments.md
@@ -0,0 +1,56 @@
+# Database query comments with Marginalia
+
+The [Marginalia gem](https://github.com/basecamp/marginalia) is used to add
+query comments containing application related context information to PostgreSQL
+queries generated by ActiveRecord.
+
+It is very useful for tracing problematic queries back to the application source.
+
+A DB Engineer during an on-call incident will have the full context of a query
+and its application source from the comments.
+
+## Metadata information in comments
+
+Queries generated from **Rails** include the following metadata in comments:
+
+- `application`
+- `controller`
+- `action`
+- `correlation_id`
+- `line`
+
+Queries generated from **Sidekiq** workers will include the following metadata
+in comments:
+
+- `application`
+- `jid`
+- `job_class`
+- `correlation_id`
+- `line`
+
+Examples of queries with comments as observed in `development.log`:
+
+1. Rails:
+
+ ```sql
+ SELECT "project_features".* FROM "project_features" WHERE "project_features"."project_id" = $1 LIMIT $2 [["project_id", 5], ["LIMIT", 1]] /*application:web,controller:jobs,action:trace,correlation_id:rYF4mey9CH3,line:/app/policies/project_policy.rb:504:in `feature_available?'*/
+ ```
+
+1. Sidekiq:
+
+ ```sql
+ SELECT "ci_pipelines".* FROM "ci_pipelines" WHERE "ci_pipelines"."id" = $1 LIMIT $2 [["id", 64], ["LIMIT", 1]] /*application:sidekiq,jid:e7d6668a39a991e323009833,job_class:ExpireJobCacheWorker,correlation_id:rYF4mey9CH3,line:/app/workers/expire_job_cache_worker.rb:14:in `perform'*/
+ ```
+
+## Enable/Disable the feature
+
+Enabling or disabling the feature requires a **restart/SIGHUP** of the Web and
+Sidekiq workers, as the feature flag's state is memoized upon starting up.
+
+The `feature_flag` for this feature is **disabled** by default. You can enable
+or disable it with:
+
+```ruby
+Feature.enable(:marginalia)
+Feature.disable(:marginalia)
+```
diff --git a/ee/changelogs/unreleased/marginalia-feature-check-on-app-start.yml b/ee/changelogs/unreleased/marginalia-feature-check-on-app-start.yml
new file mode 100644
index 0000000000000000000000000000000000000000..6715eb18f36e417412100882754fe307aa5ee0a3
--- /dev/null
+++ b/ee/changelogs/unreleased/marginalia-feature-check-on-app-start.yml
@@ -0,0 +1,5 @@
+---
+title: Added Marginalia feature which can generate PostgreSQL query comments to Gitlab
+merge_request: 21364
+author: BalaKumar
+type: added
diff --git a/lib/gitlab/marginalia.rb b/lib/gitlab/marginalia.rb
new file mode 100644
index 0000000000000000000000000000000000000000..2be96cecae31332dd1df4f3e1a4bf38b64c5b0c5
--- /dev/null
+++ b/lib/gitlab/marginalia.rb
@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module Gitlab
+ module Marginalia
+ MARGINALIA_FEATURE_FLAG = :marginalia
+
+ def self.set_application_name
+ ::Marginalia.application_name = Gitlab.process_name
+ end
+
+ def self.enable_sidekiq_instrumentation
+ if Sidekiq.server?
+ ::Marginalia::SidekiqInstrumentation.enable!
+ end
+ end
+
+ def self.cached_feature_enabled?
+ !!@enabled
+ end
+
+ def self.set_feature_cache
+ # During db:create and db:bootstrap skip feature query as DB is not available yet.
+ return false unless ActiveRecord::Base.connected? && Gitlab::Database.cached_table_exists?('features')
+
+ @enabled = Feature.enabled?(MARGINALIA_FEATURE_FLAG)
+ end
+ end
+end
diff --git a/lib/gitlab/marginalia/active_record_instrumentation.rb b/lib/gitlab/marginalia/active_record_instrumentation.rb
new file mode 100644
index 0000000000000000000000000000000000000000..3266b9f833623272190d5ae8814325be2becf0ab
--- /dev/null
+++ b/lib/gitlab/marginalia/active_record_instrumentation.rb
@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+# Patch to annotate sql only when the feature is enabled.
+module Gitlab
+ module Marginalia
+ module ActiveRecordInstrumentation
+ def annotate_sql(sql)
+ Gitlab::Marginalia.cached_feature_enabled? ? super(sql) : sql
+ end
+ end
+ end
+end
diff --git a/lib/gitlab/marginalia/comment.rb b/lib/gitlab/marginalia/comment.rb
new file mode 100644
index 0000000000000000000000000000000000000000..a0eee82376308573b9deb0d8bc3aeee548ab0462
--- /dev/null
+++ b/lib/gitlab/marginalia/comment.rb
@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+# Module to support correlation_id and additional job details.
+module Gitlab
+ module Marginalia
+ module Comment
+ private
+
+ def jid
+ bg_job["jid"] if bg_job.present?
+ end
+
+ def job_class
+ bg_job["class"] if bg_job.present?
+ end
+
+ def correlation_id
+ if bg_job.present?
+ bg_job["correlation_id"]
+ else
+ Labkit::Correlation::CorrelationId.current_id
+ end
+ end
+
+ def bg_job
+ job = ::Marginalia::Comment.marginalia_job
+
+ # We are using 'Marginalia::SidekiqInstrumentation' which does not support 'ActiveJob::Base'.
+ # Gitlab also uses 'ActionMailer::DeliveryJob' which inherits from ActiveJob::Base.
+ # So below condition is used to return metadata for such jobs.
+ if job && job.is_a?(ActionMailer::DeliveryJob)
+ {
+ "class" => job.arguments.first,
+ "jid" => job.job_id
+ }
+ else
+ job
+ end
+ end
+ end
+ end
+end
diff --git a/spec/lib/marginalia_spec.rb b/spec/lib/marginalia_spec.rb
new file mode 100644
index 0000000000000000000000000000000000000000..5dc54af99cef1524f32d58c44334eadc7102d1eb
--- /dev/null
+++ b/spec/lib/marginalia_spec.rb
@@ -0,0 +1,173 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe 'Marginalia spec' do
+ class MarginaliaTestController < ActionController::Base
+ def first_user
+ User.first
+ render body: nil
+ end
+ end
+
+ class MarginaliaTestJob
+ include Sidekiq::Worker
+
+ def perform
+ User.first
+ end
+ end
+
+ class MarginaliaTestMailer < BaseMailer
+ def first_user
+ User.first
+ end
+ end
+
+ def add_sidekiq_middleware
+ # Reference: https://github.com/mperham/sidekiq/wiki/Testing#testing-server-middlewaresidekiq
+ # Sidekiq test harness fakes worker without its server middlewares, so include instrumentation to 'Sidekiq::Testing' server middleware.
+ Sidekiq::Testing.server_middleware do |chain|
+ chain.add Marginalia::SidekiqInstrumentation::Middleware
+ end
+ end
+
+ def remove_sidekiq_middleware
+ Sidekiq::Testing.server_middleware do |chain|
+ chain.remove Marginalia::SidekiqInstrumentation::Middleware
+ end
+ end
+
+ def stub_feature(value)
+ allow(Gitlab::Marginalia).to receive(:cached_feature_enabled?).and_return(value)
+ end
+
+ def make_request(correlation_id)
+ request_env = Rack::MockRequest.env_for('/')
+
+ ::Labkit::Correlation::CorrelationId.use_id(correlation_id) do
+ MarginaliaTestController.action(:first_user).call(request_env)
+ end
+ end
+
+ describe 'For rails web requests' do
+ let(:correlation_id) { SecureRandom.uuid }
+ let(:recorded) { ActiveRecord::QueryRecorder.new { make_request(correlation_id) } }
+
+ let(:component_map) do
+ {
+ "application" => "test",
+ "controller" => "marginalia_test",
+ "action" => "first_user",
+ "line" => "/spec/support/helpers/query_recorder.rb",
+ "correlation_id" => correlation_id
+ }
+ end
+
+ context 'when the feature is enabled' do
+ before do
+ stub_feature(true)
+ end
+
+ it 'generates a query that includes the component and value' do
+ component_map.each do |component, value|
+ expect(recorded.log.last).to include("#{component}:#{value}")
+ end
+ end
+ end
+
+ context 'when the feature is disabled' do
+ before do
+ stub_feature(false)
+ end
+
+ it 'excludes annotations in generated queries' do
+ expect(recorded.log.last).not_to include("/*")
+ expect(recorded.log.last).not_to include("*/")
+ end
+ end
+ end
+
+ describe 'for Sidekiq worker jobs' do
+ before(:all) do
+ add_sidekiq_middleware
+
+ # Because of faking, 'Sidekiq.server?' does not work so implicitly set application name which is done in config/initializers/0_marginalia.rb
+ Marginalia.application_name = "sidekiq"
+ end
+
+ after(:all) do
+ MarginaliaTestJob.clear
+ remove_sidekiq_middleware
+ end
+
+ around do |example|
+ Sidekiq::Testing.fake! { example.run }
+ end
+
+ before do
+ MarginaliaTestJob.perform_async
+ end
+
+ let(:sidekiq_job) { MarginaliaTestJob.jobs.first }
+ let(:recorded) { ActiveRecord::QueryRecorder.new { MarginaliaTestJob.drain } }
+
+ let(:component_map) do
+ {
+ "application" => "sidekiq",
+ "job_class" => "MarginaliaTestJob",
+ "line" => "/spec/support/sidekiq_middleware.rb",
+ "correlation_id" => sidekiq_job['correlation_id'],
+ "jid" => sidekiq_job['jid']
+ }
+ end
+
+ context 'when the feature is enabled' do
+ before do
+ stub_feature(true)
+ end
+
+ it 'generates a query that includes the component and value' do
+ component_map.each do |component, value|
+ expect(recorded.log.last).to include("#{component}:#{value}")
+ end
+ end
+
+ describe 'for ActionMailer delivery jobs' do
+ let(:delivery_job) { MarginaliaTestMailer.first_user.deliver_later }
+
+ let(:recorded) do
+ ActiveRecord::QueryRecorder.new do
+ delivery_job.perform_now
+ end
+ end
+
+ let(:component_map) do
+ {
+ "application" => "sidekiq",
+ "line" => "/lib/gitlab/i18n.rb",
+ "jid" => delivery_job.job_id,
+ "job_class" => delivery_job.arguments.first
+ }
+ end
+
+ it 'generates a query that includes the component and value' do
+ component_map.each do |component, value|
+ expect(recorded.log.last).to include("#{component}:#{value}")
+ end
+ end
+ end
+ end
+
+ context 'when the feature is disabled' do
+ before do
+ stub_feature(false)
+ end
+
+ it 'excludes annotations in generated queries' do
+ expect(recorded.log.last).not_to include("/*")
+ expect(recorded.log.last).not_to include("*/")
+ end
+ end
+ end
+end
diff --git a/spec/serializers/pipeline_serializer_spec.rb b/spec/serializers/pipeline_serializer_spec.rb
index 7661c8acc135eb4b17ccf4f7113870ac91c86cd4..f1f761a6fd0232518c7981c667bc1b378427b84f 100644
--- a/spec/serializers/pipeline_serializer_spec.rb
+++ b/spec/serializers/pipeline_serializer_spec.rb
@@ -130,10 +130,10 @@
it 'preloads related merge requests' do
recorded = ActiveRecord::QueryRecorder.new { subject }
+ expected_query = "SELECT \"merge_requests\".* FROM \"merge_requests\" " \
+ "WHERE \"merge_requests\".\"id\" IN (#{merge_request_1.id}, #{merge_request_2.id})"
- expect(recorded.log)
- .to include("SELECT \"merge_requests\".* FROM \"merge_requests\" " \
- "WHERE \"merge_requests\".\"id\" IN (#{merge_request_1.id}, #{merge_request_2.id})")
+ expect(recorded.log).to include(a_string_starting_with(expected_query))
end
end
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index 2afe80f455ebba492fbff32b3ff01dd6981594ed..1f0119108a8504128bba602dbc93c392b4483552 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -159,6 +159,9 @@
.with(:force_autodevops_on_by_default, anything)
.and_return(false)
+ # Enable Marginalia feature for all specs in the test suite.
+ allow(Gitlab::Marginalia).to receive(:cached_feature_enabled?).and_return(true)
+
# The following can be removed once Vue Issuable Sidebar
# is feature-complete and can be made default in place
# of older sidebar.