Set user to the current scope via sidekiq middleware (#1469)

shouichi · web-flow · commit 79f04538d676 · 2021-06-25T11:58:10.000+08:00
diff --git a/sentry-sidekiq/lib/sentry-sidekiq.rb b/sentry-sidekiq/lib/sentry-sidekiq.rb
@@ -26,7 +26,12 @@ class Railtie < ::Rails::Railtie
 Sidekiq.configure_server do |config|
   config.error_handlers << Sentry::Sidekiq::ErrorHandler.new
   config.server_middleware do |chain|
-    chain.add Sentry::Sidekiq::SentryContextMiddleware
+    chain.add Sentry::Sidekiq::SentryContextServerMiddleware
   end
 end
 
+Sidekiq.configure_client do |config|
+  config.client_middleware do |chain|
+    chain.add Sentry::Sidekiq::SentryContextClientMiddleware
+  end
+end
diff --git a/sentry-sidekiq/lib/sentry/sidekiq/sentry_context_middleware.rb b/sentry-sidekiq/lib/sentry/sidekiq/sentry_context_middleware.rb
@@ -2,14 +2,17 @@
 
 module Sentry
   module Sidekiq
-    class SentryContextMiddleware
+    class SentryContextServerMiddleware
       def call(_worker, job, queue)
         return yield unless Sentry.initialized?
 
         context_filter = Sentry::Sidekiq::ContextFilter.new(job)
 
         Sentry.clone_hub_to_current_thread
         scope = Sentry.get_current_scope
+        if (user = job["sentry_user"])
+          scope.set_user(user)
+        end
         scope.set_tags(queue: queue, jid: job["jid"])
         scope.set_contexts(sidekiq: job.merge("queue" => queue))
         scope.set_transaction_name(context_filter.transaction_name)
@@ -36,5 +39,15 @@ def finish_transaction(transaction, status)
         transaction.finish
       end
     end
+
+    class SentryContextClientMiddleware
+      def call(_worker_class, job, _queue, _redis_pool)
+        return yield unless Sentry.initialized?
+
+        user = Sentry.get_current_scope.user
+        job["sentry_user"] = user unless user.empty?
+        yield
+      end
+    end
   end
 end
diff --git a/sentry-sidekiq/spec/sentry/sidekiq/sentry_context_middleware_spec.rb b/sentry-sidekiq/spec/sentry/sidekiq/sentry_context_middleware_spec.rb
@@ -0,0 +1,109 @@
+require "spec_helper"
+
+RSpec.shared_context "sidekiq", shared_context: :metadata do
+  let(:user) { { "id" => rand(10_000) } }
+
+  let(:client) do
+    Sidekiq::Client.new.tap do |client|
+      client.middleware do |chain|
+        chain.add Sentry::Sidekiq::SentryContextClientMiddleware
+      end
+    end
+  end
+
+  let(:random_empty_queue) do
+    Sidekiq::Queue.new(rand(10_000)).tap do |queue|
+      queue.clear
+    end
+  end
+end
+
+RSpec.describe Sentry::Sidekiq::SentryContextServerMiddleware do
+  include_context "sidekiq"
+
+  after { Sidekiq::RetrySet.new.clear }
+
+  it "sets user from the job to events" do
+    perform_basic_setup { |config| config.traces_sample_rate = 1.0 }
+    Sentry.set_user(user)
+
+    queue = random_empty_queue
+    options = { queues: [queue.name] }
+    processor = Sidekiq::Manager.new(options).workers.first
+
+    client.push('queue' => queue.name, 'class' => HappyWorker, 'args' => [])
+
+    expect { processor.send(:process_one) }.
+      to change { Sentry.get_current_client.transport.events.size }.by(1)
+
+    event =  Sentry.get_current_client.transport.events.first
+    expect(event).not_to be_nil
+    expect(event.user).to eq(user)
+  end
+
+  it "sets user from the job to sidekiq event if worker raises an exception" do
+    perform_basic_setup { |config| config.traces_sample_rate = 0 }
+    Sentry.set_user(user)
+
+    queue = random_empty_queue
+    options = { queues: [queue.name] }
+    processor = Sidekiq::Manager.new(options).workers.first
+
+    client.push('queue' => queue.name, 'class' => SadWorker, 'args' => [])
+
+    expect do
+      begin; processor.send(:process_one); rescue RuntimeError; end
+    end.
+      to change { Sentry.get_current_client.transport.events.size }.by(1)
+
+    Sentry.get_current_client.transport.events.each do |event|
+      expect(event.user).to eq(user)
+    end
+  end
+
+  it "sets user from the job to sidekiq and error events if worker raises an exception when trace enabled" do
+    perform_basic_setup { |config| config.traces_sample_rate = 1.0 }
+    Sentry.set_user(user)
+
+    queue = random_empty_queue
+    options = { queues: [queue.name] }
+    processor = Sidekiq::Manager.new(options).workers.first
+
+    client.push('queue' => queue.name, 'class' => SadWorker, 'args' => [])
+
+    # XXX: In ruby 2.4, two events are pushed. In other versions, only one
+    # event is pushed. Use by_at_least.
+    expect do
+      begin; processor.send(:process_one); rescue RuntimeError; end
+    end.
+      to change { Sentry.get_current_client.transport.events.size }.by_at_least(1)
+
+    Sentry.get_current_client.transport.events.each do |event|
+      expect(event.user).to eq(user)
+    end
+  end
+end
+
+RSpec.describe Sentry::Sidekiq::SentryContextClientMiddleware do
+  include_context "sidekiq"
+
+  before { perform_basic_setup }
+
+  it "does not user to the job if user is absence in the current scope" do
+    queue = random_empty_queue
+    client.push('queue' => queue.name, 'class' => HappyWorker, 'args' => [])
+
+    expect(queue.size).to be(1)
+    expect(queue.first["sentry_user"]).to be_nil
+  end
+
+  it "sets user of the current scope to the job if present" do
+    queue = random_empty_queue
+    Sentry.set_user(user)
+
+    client.push('queue' => queue.name, 'class' => HappyWorker, 'args' => [])
+
+    expect(queue.size).to be(1)
+    expect(queue.first["sentry_user"]).to eq(user)
+  end
+end
diff --git a/sentry-sidekiq/spec/sentry/sidekiq_spec.rb b/sentry-sidekiq/spec/sentry/sidekiq_spec.rb
@@ -24,7 +24,7 @@
 
   it "registers error handlers and middlewares" do
     expect(Sidekiq.error_handlers).to include(described_class::ErrorHandler)
-    expect(Sidekiq.server_middleware.entries.first.klass).to eq(described_class::SentryContextMiddleware)
+    expect(Sidekiq.server_middleware.entries.first.klass).to eq(described_class::SentryContextServerMiddleware)
   end
 
   it "captues exception raised in the worker" do
