diff --git a/lib/sidekiq/middleware/current_attributes.rb b/lib/sidekiq/middleware/current_attributes.rb new file mode 100644 index 00000000..44833882 --- /dev/null +++ b/lib/sidekiq/middleware/current_attributes.rb @@ -0,0 +1,48 @@ +require "active_support/current_attributes" + +module Sidekiq + ## + # Automatically save and load any current attributes in the execution context + # so context attributes "flow" from Rails actions into any associated jobs. + # This can be useful for multi-tenancy, i18n locale, timezone, any implicit + # per-request attribute. See +ActiveSupport::CurrentAttributes+. + # + # @example + # + # # in your initializer + # require "sidekiq/middleware/current_attributes" + # Sidekiq::CurrentAttributes.persist(Myapp::Current) + # + module CurrentAttributes + class Save + def initialize(with:) + @klass = with + end + + def call(_, job, _, _) + job["ctx"] = @klass.attributes + yield + end + end + + class Load + def initialize(with:) + @klass = with + end + + def call(_, job, _, &block) + @klass.set(job["ctx"], &block) + end + end + + def self.persist(klass) + Sidekiq.configure_client do |config| + config.client_middleware.add Save, with: klass + end + Sidekiq.configure_server do |config| + config.client_middleware.add Save, with: klass + config.server_middleware.add Load, with: klass + end + end + end +end diff --git a/test/test_current_attributes.rb b/test/test_current_attributes.rb new file mode 100644 index 00000000..7f0aa2ae --- /dev/null +++ b/test/test_current_attributes.rb @@ -0,0 +1,51 @@ +require_relative "./helper" +require "sidekiq/middleware/current_attributes" + +module Myapp + class Current < ActiveSupport::CurrentAttributes + attribute :user_id + end +end + +class TestCurrentAttributes < Minitest::Test + def test_save + cm = Sidekiq::CurrentAttributes::Save.new(with: Myapp::Current) + job = {} + with_context(:user_id, 123) do + cm.call(nil, job, nil, nil) do + assert_equal 123, job["ctx"][:user_id] + end + end + end + + def test_load + cm = Sidekiq::CurrentAttributes::Load.new(with: Myapp::Current) + + job = { "ctx" => { "user_id" => 123 } } + assert_nil Myapp::Current.user_id + cm.call(nil, job, nil) do + assert_equal 123, Myapp::Current.user_id + end + # the Rails reloader is responsible for reseting Current after every unit of work + end + + def test_persist + begin + Sidekiq::CurrentAttributes.persist(Myapp::Current) + ensure + Sidekiq.client_middleware.clear + Sidekiq.server_middleware.clear + end + end + + private + + def with_context(attr, value) + begin + Myapp::Current.send("#{attr}=", value) + yield + ensure + Myapp::Current.reset_all + end + end +end