From 6ea85aff7169c528b47704be4370d7043b52f7af Mon Sep 17 00:00:00 2001
From: Jean Boussier <jean.boussier@gmail.com>
Date: Mon, 23 Aug 2021 15:37:14 +0200
Subject: [PATCH] Make `QueryLogs.set_context` restore previous values

It's really the least surpising behavior for block based APIs like this one
and is consistent with `with_tags`.
---
 activerecord/lib/active_record/query_logs.rb |  8 +++++---
 activerecord/test/cases/query_logs_test.rb   | 14 ++++++++++++++
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/activerecord/lib/active_record/query_logs.rb b/activerecord/lib/active_record/query_logs.rb
index c37b29418ca36..ef489a6a12e1c 100644
--- a/activerecord/lib/active_record/query_logs.rb
+++ b/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.
diff --git a/activerecord/test/cases/query_logs_test.rb b/activerecord/test/cases/query_logs_test.rb
index 2787099965b03..e3eaa60c9aa67 100644
--- a/activerecord/test/cases/query_logs_test.rb
+++ b/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