Support multiple current attributes in middleware

lib/sidekiq/middleware/current_attributes.rb

@@ -7,26 +7,32 @@ module Sidekiq
   # This can be useful for multi-tenancy, i18n locale, timezone, any implicit
   # per-request attribute. See +ActiveSupport::CurrentAttributes+.
   #
+  # For multiple current attributes, pass an array of current attributes.
+  #
   # @example
   #
   #   # in your initializer
   #   require "sidekiq/middleware/current_attributes"
   #   Sidekiq::CurrentAttributes.persist("Myapp::Current")
+  #   # or multiple current attributes
+  #   Sidekiq::CurrentAttributes.persist(["Myapp::Current", "Myapp::OtherCurrent"])
   #
   module CurrentAttributes
     class Save
       include Sidekiq::ClientMiddleware
 
-      def initialize(cattr)
-        @strklass = cattr
+      def initialize(cattrs)
+        @cattrs = cattrs
       end
 
       def call(_, job, _, _)
-        if !job.has_key?("cattr")
-          attrs = @strklass.constantize.attributes
-          # Retries can push the job N times, we don't
-          # want retries to reset cattr. #5692, #5090
-          job["cattr"] = attrs if attrs.any?
+        @cattrs.each do |(key, strklass)|
+          if !job.has_key?(key)
+            attrs = strklass.constantize.attributes
+            # Retries can push the job N times, we don't
+            # want retries to reset cattr. #5692, #5090
+            job[key] = attrs if attrs.any?
+          end
         end
         yield
       end
@@ -35,22 +41,55 @@ def call(_, job, _, _)
     class Load
       include Sidekiq::ServerMiddleware
 
-      def initialize(cattr)
-        @strklass = cattr
+      def initialize(cattrs)
+        @cattrs = cattrs
       end
 
       def call(_, job, _, &block)
-        if job.has_key?("cattr")
-          @strklass.constantize.set(job["cattr"], &block)
-        else
-          yield
+        cattrs_to_reset = []
+
+        @cattrs.each do |(key, strklass)|
+          if job.has_key?(key)
+            constklass = strklass.constantize
+            cattrs_to_reset << constklass
+
+            job[key].each do |(attribute, value)|
+              constklass.public_send("#{attribute}=", value)
+            end
+          end
         end
+
+        yield
+      ensure
+        cattrs_to_reset.each(&:reset)
       end
     end
 
-    def self.persist(klass, config = Sidekiq.default_configuration)
-      config.client_middleware.add Save, klass.to_s
-      config.server_middleware.add Load, klass.to_s
+    class << self
+      def persist(klass_or_array, config = Sidekiq.default_configuration)
+        cattrs = build_cattrs_hash(klass_or_array)
+
+        config.client_middleware.add Save, cattrs
+        config.server_middleware.add Load, cattrs
+      end
+
+      private
+
+      def build_cattrs_hash(klass_or_array)
+        if klass_or_array.is_a?(Array)
+          Hash.new.tap do |hash|
+            klass_or_array.each_with_index do |klass, index|
+              hash[key_at(index)] = klass.to_s
+            end
+          end
+        else
+          {key_at(0) => klass_or_array.to_s}
+        end
+      end
+
+      def key_at(index)
+        index == 0 ? "cattr" : "cattr_#{index}"
+      end
     end
   end
 end

test/current_attributes_test.rb

@@ -8,6 +8,10 @@ module Myapp
   class Current < ActiveSupport::CurrentAttributes
     attribute :user_id
   end
+
+  class OtherCurrent < ActiveSupport::CurrentAttributes
+    attribute :other_id
+  end
 end
 
 describe "Current attributes" do
@@ -16,36 +20,50 @@ class Current < ActiveSupport::CurrentAttributes
   end
 
   it "saves" do
-    cm = Sidekiq::CurrentAttributes::Save.new("Myapp::Current")
+    cm = Sidekiq::CurrentAttributes::Save.new({
+      "cattr" => "Myapp::Current",
+      "cattr_1" => "Myapp::OtherCurrent"
+    })
     job = {}
-    with_context(:user_id, 123) do
-      cm.call(nil, job, nil, nil) do
-        assert_equal 123, job["cattr"][:user_id]
+    with_context("Myapp::Current", :user_id, 123) do
+      with_context("Myapp::OtherCurrent", :other_id, 789) do
+        cm.call(nil, job, nil, nil) do
+          assert_equal 123, job["cattr"][:user_id]
+          assert_equal 789, job["cattr_1"][:other_id]
+        end
       end
     end
 
-    with_context(:user_id, 456) do
-      cm.call(nil, job, nil, nil) do
-        assert_equal 123, job["cattr"][:user_id]
+    with_context("Myapp::Current", :user_id, 456) do
+      with_context("Myapp::OtherCurrent", :other_id, 999) do
+        cm.call(nil, job, nil, nil) do
+          assert_equal 123, job["cattr"][:user_id]
+          assert_equal 789, job["cattr_1"][:other_id]
+        end
       end
     end
   end
 
   it "loads" do
-    cm = Sidekiq::CurrentAttributes::Load.new("Myapp::Current")
+    cm = Sidekiq::CurrentAttributes::Load.new({
+      "cattr" => "Myapp::Current",
+      "cattr_1" => "Myapp::OtherCurrent"
+    })
 
-    job = {"cattr" => {"user_id" => 123}}
+    job = {"cattr" => {"user_id" => 123}, "cattr_1" => {"other_id" => 456}}
     assert_nil Myapp::Current.user_id
+    assert_nil Myapp::OtherCurrent.other_id
     cm.call(nil, job, nil) do
       assert_equal 123, Myapp::Current.user_id
+      assert_equal 456, Myapp::OtherCurrent.other_id
     end
     # the Rails reloader is responsible for reseting Current after every unit of work
   end
 
-  it "persists" do
+  it "persists with class argument" do
     Sidekiq::CurrentAttributes.persist("Myapp::Current", @config)
     job_hash = {}
-    with_context(:user_id, 16) do
+    with_context("Myapp::Current", :user_id, 16) do
       @config.client_middleware.invoke(nil, job_hash, nil, nil) do
         assert_equal 16, job_hash["cattr"][:user_id]
       end
@@ -62,12 +80,27 @@ class Current < ActiveSupport::CurrentAttributes
     #   Sidekiq.server_middleware.clear
   end
 
+  it "persists with hash argument" do
+    cattrs = [Myapp::Current, "Myapp::OtherCurrent"]
+    Sidekiq::CurrentAttributes.persist(cattrs, @config)
+    job_hash = {}
+    with_context("Myapp::Current", :user_id, 16) do
+      with_context("Myapp::OtherCurrent", :other_id, 17) do
+        @config.client_middleware.invoke(nil, job_hash, nil, nil) do
+          assert_equal 16, job_hash["cattr"][:user_id]
+          assert_equal 17, job_hash["cattr_1"][:other_id]
+        end
+      end
+    end
+  end
+
   private
 
-  def with_context(attr, value)
-    Myapp::Current.send("#{attr}=", value)
+  def with_context(strklass, attr, value)
+    constklass = strklass.constantize
+    constklass.send("#{attr}=", value)
     yield
   ensure
-    Myapp::Current.reset_all
+    constklass.reset_all
   end
 end