Merge pull request rails#43076 from Shopify/ar-query-logs-nesting

Make `QueryLogs.set_context` restore previous values

Author: byroot

activerecord/lib/active_record/query_logs.rb

@@ -96,13 +96,15 @@ def update_context(**options)
       end
 
       # Updates the context used to construct tags in the SQL comment during
-      # execution of the provided block. Resets provided values to nil after
-      # the block is executed.
+      # execution of the provided block. Resets the provided keys to their
+      # previous value once the block exits.
       def set_context(**options)
+        keys = options.keys
+        previous_context = keys.zip(context.values_at(*keys)).to_h
         update_context(**options)
         yield if block_given?
       ensure
-        update_context(**options.transform_values! { NullObject.new })
+        update_context(**previous_context)
       end
 
       # Temporarily tag any query executed within `&block`. Can be nested.

activerecord/test/cases/query_logs_test.rb

@@ -253,4 +253,18 @@ def test_custom_proc_context_tags
     ActiveRecord::QueryLogs.update_context(foo: nil)
     ActiveRecord::QueryLogs.tags = original_tags
   end
+
+  def test_set_context_restore_state
+    original_tags = ActiveRecord::QueryLogs.tags
+    ActiveRecord::QueryLogs.tags = [foo: -> { context[:foo] }]
+    ActiveRecord::QueryLogs.set_context(foo: "bar") do
+      assert_sql(%r{/\*foo:bar\*/$}) { Dashboard.first }
+      ActiveRecord::QueryLogs.set_context(foo: "plop") do
+        assert_sql(%r{/\*foo:plop\*/$}) { Dashboard.first }
+      end
+      assert_sql(%r{/\*foo:bar\*/$}) { Dashboard.first }
+    end
+  ensure
+    ActiveRecord::QueryLogs.tags = original_tags
+  end
 end