Add support for Active Record query log tags (#33342)

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'

lib/mastodon/sidekiq_middleware.rb

@@ -3,8 +3,10 @@
 class Mastodon::SidekiqMiddleware
   BACKTRACE_LIMIT = 3
 
-  def call(*, &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