From 3f8f06eb466e00fb7db3fc84dba2d820b9286993 Mon Sep 17 00:00:00 2001
From: Renaud Chaput <renchap@gmail.com>
Date: Wed, 18 Dec 2024 11:05:00 +0100
Subject: [PATCH] Add support for Active Record query log tags (#33342)

---
 config/application.rb              |  5 +++++
 lib/mastodon/sidekiq_middleware.rb | 16 ++++++++++++++--
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/config/application.rb b/config/application.rb
index e4e9680e66..8939ad91f4 100644
--- a/config/application.rb
+++ b/config/application.rb
@@ -108,6 +108,11 @@ module Mastodon
     config.x.mastodon = config_for(:mastodon)
     config.x.translation = config_for(:translation)
 
+    if ENV.fetch('QUERY_LOG_TAGS_ENABLED', 'false') == 'true'
+      config.active_record.query_log_tags_enabled = ENV.fetch('QUERY_LOG_TAGS_ENABLED', 'false') == 'true'
+      config.active_record.query_log_tags = [:namespaced_controller, :action, :sidekiq_job_class]
+    end
+
     config.to_prepare do
       Doorkeeper::AuthorizationsController.layout 'modal'
       Doorkeeper::AuthorizedApplicationsController.layout 'admin'
diff --git a/lib/mastodon/sidekiq_middleware.rb b/lib/mastodon/sidekiq_middleware.rb
index 8ce1124c69..4df22002c8 100644
--- a/lib/mastodon/sidekiq_middleware.rb
+++ b/lib/mastodon/sidekiq_middleware.rb
@@ -3,8 +3,10 @@
 class Mastodon::SidekiqMiddleware
   BACKTRACE_LIMIT = 3
 
-  def call(*, &block)
-    Chewy.strategy(:mastodon, &block)
+  def call(_worker_class, job, _queue, &block)
+    setup_query_log_tags(job) do
+      Chewy.strategy(:mastodon, &block)
+    end
   rescue Mastodon::HostValidationError
     # Do not retry
   rescue => e
@@ -61,4 +63,14 @@ class Mastodon::SidekiqMiddleware
     Thread.current[:statsd_socket]&.close
     Thread.current[:statsd_socket] = nil
   end
+
+  def setup_query_log_tags(job, &block)
+    if Rails.configuration.active_record.query_log_tags_enabled
+      # If `wrapped` is set, this is an `ActiveJob` which is already in the execution context
+      sidekiq_job_class = job['wrapped'].present? ? nil : job['class'].to_s
+      ActiveSupport::ExecutionContext.set(sidekiq_job_class: sidekiq_job_class, &block)
+    else
+      yield
+    end
+  end
 end