From 6bad959565900499325d4239266cbb1f1f8e4056 Mon Sep 17 00:00:00 2001
From: Jean Boussier <jean.boussier@gmail.com>
Date: Fri, 5 Nov 2021 13:09:14 +0100
Subject: [PATCH] Extract ActiveSupport::ExecutionContext out of
 ActiveRecord::QueryLogs

I'm working on a standardized error reporting interface
(https://github.com/rails/rails/issues/43472) and it has the same need
for a `context` than Active Record's query logs.

Just like query logs, when reporting an error you want to know what the
current controller or job is, etc.

So by extracting it we can allow both API to use the same store.
---
 actionpack/lib/abstract_controller/base.rb    |  1 +
 actionpack/lib/action_controller.rb           |  1 -
 .../lib/action_controller/metal/query_tags.rb | 16 -----
 actionpack/lib/action_controller/railtie.rb   |  4 --
 activejob/lib/active_job/execution.rb         |  1 +
 activejob/lib/active_job/query_tags.rb        | 16 -----
 activejob/lib/active_job/railtie.rb           |  4 --
 activerecord/lib/active_record/query_logs.rb  | 42 +++--------
 activerecord/lib/active_record/railtie.rb     |  4 --
 activerecord/test/cases/query_logs_test.rb    | 71 ++++---------------
 activesupport/lib/active_support.rb           |  1 +
 .../lib/active_support/execution_context.rb   | 53 ++++++++++++++
 .../execution_context/test_helper.rb          | 13 ++++
 activesupport/lib/active_support/railtie.rb   |  9 +++
 activesupport/test/current_attributes_test.rb |  9 +--
 activesupport/test/execution_context_test.rb  | 47 ++++++++++++
 16 files changed, 151 insertions(+), 141 deletions(-)
 delete mode 100644 actionpack/lib/action_controller/metal/query_tags.rb
 delete mode 100644 activejob/lib/active_job/query_tags.rb
 create mode 100644 activesupport/lib/active_support/execution_context.rb
 create mode 100644 activesupport/lib/active_support/execution_context/test_helper.rb
 create mode 100644 activesupport/test/execution_context_test.rb

diff --git a/actionpack/lib/abstract_controller/base.rb b/actionpack/lib/abstract_controller/base.rb
index 61fc35561e..11403d8f7d 100644
--- a/actionpack/lib/abstract_controller/base.rb
+++ b/actionpack/lib/abstract_controller/base.rb
@@ -148,6 +148,7 @@ module AbstractController
 
       @_response_body = nil
 
+      ActiveSupport::ExecutionContext[:controller] = self
       process_action(action_name, *args)
     end
 
diff --git a/actionpack/lib/action_controller.rb b/actionpack/lib/action_controller.rb
index d80b4f5b03..62ae8f8d9e 100644
--- a/actionpack/lib/action_controller.rb
+++ b/actionpack/lib/action_controller.rb
@@ -37,7 +37,6 @@ module ActionController
     autoload :Logging
     autoload :MimeResponds
     autoload :ParamsWrapper
-    autoload :QueryTags
     autoload :Redirecting
     autoload :Renderers
     autoload :Rendering
diff --git a/actionpack/lib/action_controller/metal/query_tags.rb b/actionpack/lib/action_controller/metal/query_tags.rb
deleted file mode 100644
index f8ca992f0f..0000000000
--- a/actionpack/lib/action_controller/metal/query_tags.rb
+++ /dev/null
@@ -1,16 +0,0 @@
-# frozen_string_literal: true
-
-module ActionController
-  module QueryTags # :nodoc:
-    extend ActiveSupport::Concern
-
-    included do
-      around_action :expose_controller_to_query_logs
-    end
-
-    private
-      def expose_controller_to_query_logs(&block)
-        ActiveRecord::QueryLogs.set_context(controller: self, &block)
-      end
-  end
-end
diff --git a/actionpack/lib/action_controller/railtie.rb b/actionpack/lib/action_controller/railtie.rb
index bea6873a50..701ddb5bef 100644
--- a/actionpack/lib/action_controller/railtie.rb
+++ b/actionpack/lib/action_controller/railtie.rb
@@ -113,10 +113,6 @@ module ActionController
       if query_logs_tags_enabled
         app.config.active_record.query_log_tags += [:controller, :action]
 
-        ActiveSupport.on_load(:action_controller) do
-          include ActionController::QueryTags
-        end
-
         ActiveSupport.on_load(:active_record) do
           ActiveRecord::QueryLogs.taggings.merge!(
             controller:            ->(context) { context[:controller]&.controller_name },
diff --git a/activejob/lib/active_job/execution.rb b/activejob/lib/active_job/execution.rb
index 6f07c02122..12c1b468c3 100644
--- a/activejob/lib/active_job/execution.rb
+++ b/activejob/lib/active_job/execution.rb
@@ -54,6 +54,7 @@ module ActiveJob
 
     private
       def _perform_job
+        ActiveSupport::ExecutionContext[:job] = self
         run_callbacks :perform do
           perform(*arguments)
         end
diff --git a/activejob/lib/active_job/query_tags.rb b/activejob/lib/active_job/query_tags.rb
deleted file mode 100644
index 08b929dca9..0000000000
--- a/activejob/lib/active_job/query_tags.rb
+++ /dev/null
@@ -1,16 +0,0 @@
-# frozen_string_literal: true
-
-module ActiveJob
-  module QueryTags # :nodoc:
-    extend ActiveSupport::Concern
-
-    included do
-      around_perform :expose_job_to_query_logs
-    end
-
-    private
-      def expose_job_to_query_logs(&block)
-        ActiveRecord::QueryLogs.set_context(job: self, &block)
-      end
-  end
-end
diff --git a/activejob/lib/active_job/railtie.rb b/activejob/lib/active_job/railtie.rb
index 7755b66e28..27fb31062c 100644
--- a/activejob/lib/active_job/railtie.rb
+++ b/activejob/lib/active_job/railtie.rb
@@ -65,10 +65,6 @@ module ActiveJob
       if query_logs_tags_enabled
         app.config.active_record.query_log_tags << :job
 
-        ActiveSupport.on_load(:active_job) do
-          include ActiveJob::QueryTags
-        end
-
         ActiveSupport.on_load(:active_record) do
           ActiveRecord::QueryLogs.taggings[:job] = ->(context) { context[:job].class.name if context[:job] }
         end
diff --git a/activerecord/lib/active_record/query_logs.rb b/activerecord/lib/active_record/query_logs.rb
index 452cdb0087..f116a154dd 100644
--- a/activerecord/lib/active_record/query_logs.rb
+++ b/activerecord/lib/active_record/query_logs.rb
@@ -44,17 +44,17 @@ module ActiveRecord
   #    ]
   #    ActiveRecord::QueryLogs.tags = tags
   #
-  # The QueryLogs +context+ can be manipulated via the +set_context+ method.
+  # The QueryLogs +context+ can be manipulated via the +ActiveSupport::ExecutionContext.set+ method.
   #
   # Temporary updates limited to the execution of a block:
   #
-  #    ActiveRecord::QueryLogs.set_context(foo: Bar.new) do
+  #    ActiveSupport::ExecutionContext.set(foo: Bar.new) do
   #      posts = Post.all
   #    end
   #
   # Direct updates to a context value:
   #
-  #    ActiveRecord::QueryLogs.set_context(foo: Bar.new)
+  #    ActiveSupport::ExecutionContext[:foo] = Bar.new
   #
   # Tag comments can be prepended to the query:
   #
@@ -76,30 +76,6 @@ module ActiveRecord
     thread_mattr_accessor :cached_comment, instance_accessor: false
 
     class << self
-      # Updates the context used to construct tags in the SQL comment.
-      # If a block is given, it resets the provided keys to their
-      # previous value once the block exits.
-      def set_context(**options)
-        options.symbolize_keys!
-
-        keys = options.keys
-        previous_context = keys.zip(context.values_at(*keys)).to_h
-        context.merge!(options)
-        self.cached_comment = nil
-        if block_given?
-          begin
-            yield
-          ensure
-            context.merge!(previous_context)
-            self.cached_comment = nil
-          end
-        end
-      end
-
-      def clear_context # :nodoc:
-        context.clear
-      end
-
       def call(sql) # :nodoc:
         if prepend_comment
           "#{self.comment} #{sql}"
@@ -108,6 +84,12 @@ module ActiveRecord
         end.strip
       end
 
+      def clear_cache # :nodoc:
+        self.cached_comment = nil
+      end
+
+      ActiveSupport::ExecutionContext.after_change { ActiveRecord::QueryLogs.clear_cache }
+
       private
         # Returns an SQL comment +String+ containing the query log tags.
         # Sets and returns a cached comment if <tt>cache_query_log_tags</tt> is +true+.
@@ -126,15 +108,13 @@ module ActiveRecord
           end
         end
 
-        def context
-          Thread.current[:active_record_query_log_tags_context] ||= {}
-        end
-
         def escape_sql_comment(content)
           content.to_s.gsub(%r{ (/ (?: | \g<1>) \*) \+? \s* | \s* (\* (?: | \g<2>) /) }x, "")
         end
 
         def tag_content
+          context = ActiveSupport::ExecutionContext.to_h
+
           tags.flat_map { |i| [*i] }.filter_map do |tag|
             key, handler = tag
             handler ||= taggings[key]
diff --git a/activerecord/lib/active_record/railtie.rb b/activerecord/lib/active_record/railtie.rb
index 950ac63eb6..04112bc2e1 100644
--- a/activerecord/lib/active_record/railtie.rb
+++ b/activerecord/lib/active_record/railtie.rb
@@ -376,10 +376,6 @@ To keep using the current cache store, you can turn off cache versioning entirel
           if app.config.active_record.cache_query_log_tags
             ActiveRecord::QueryLogs.cache_query_log_tags = true
           end
-
-          app.reloader.before_class_unload { ActiveRecord::QueryLogs.clear_context }
-          app.executor.to_run              { ActiveRecord::QueryLogs.clear_context }
-          app.executor.to_complete         { ActiveRecord::QueryLogs.clear_context }
         end
       end
     end
diff --git a/activerecord/test/cases/query_logs_test.rb b/activerecord/test/cases/query_logs_test.rb
index cc34f3a39b..05207f17e3 100644
--- a/activerecord/test/cases/query_logs_test.rb
+++ b/activerecord/test/cases/query_logs_test.rb
@@ -11,13 +11,14 @@ class QueryLogsTest < ActiveRecord::TestCase
   }
 
   def setup
-    # QueryLogs context is automatically reset in Rails app via an executor hooks set in railtie
+    # ActiveSupport::ExecutionContext context is automatically reset in Rails app via an executor hooks set in railtie
     # But not in Active Record's own test suite.
-    ActiveRecord::QueryLogs.clear_context
+    ActiveSupport::ExecutionContext.clear
 
     # Enable the query tags logging
     @original_transformers = ActiveRecord.query_transformers
     @original_prepend = ActiveRecord::QueryLogs.prepend_comment
+    @original_tags = ActiveRecord::QueryLogs.tags
     ActiveRecord.query_transformers += [ActiveRecord::QueryLogs]
     ActiveRecord::QueryLogs.prepend_comment = false
     ActiveRecord::QueryLogs.cache_query_log_tags = false
@@ -27,10 +28,13 @@ class QueryLogsTest < ActiveRecord::TestCase
   def teardown
     ActiveRecord.query_transformers = @original_transformers
     ActiveRecord::QueryLogs.prepend_comment = @original_prepend
-    ActiveRecord::QueryLogs.tags = []
-    # QueryLogs context is automatically reset in Rails app via an executor hooks set in railtie
+    ActiveRecord::QueryLogs.tags = @original_tags
+    ActiveRecord::QueryLogs.prepend_comment = false
+    ActiveRecord::QueryLogs.cache_query_log_tags = false
+    ActiveRecord::QueryLogs.cached_comment = nil
+    # ActiveSupport::ExecutionContext context is automatically reset in Rails app via an executor hooks set in railtie
     # But not in Active Record's own test suite.
-    ActiveRecord::QueryLogs.clear_context
+    ActiveSupport::ExecutionContext.clear
   end
 
   def test_escaping_good_comment
@@ -57,8 +61,6 @@ class QueryLogsTest < ActiveRecord::TestCase
     assert_sql(%r{/\*application:active_record\*/ select id from posts$}) do
       ActiveRecord::Base.connection.execute "select id from posts"
     end
-  ensure
-    ActiveRecord::QueryLogs.prepend_comment = nil
   end
 
   def test_exists_is_commented
@@ -112,41 +114,24 @@ class QueryLogsTest < ActiveRecord::TestCase
     ActiveRecord::QueryLogs.stub(:cached_comment, "/*cached_comment*/") do
       assert_equal "/*cached_comment*/", ActiveRecord::QueryLogs.call("")
     end
-  ensure
-    ActiveRecord::QueryLogs.cached_comment = nil
-    ActiveRecord::QueryLogs.cache_query_log_tags = false
   end
 
   def test_resets_cache_on_context_update
     ActiveRecord::QueryLogs.cache_query_log_tags = true
-    ActiveRecord::QueryLogs.set_context(temporary: "value")
+    ActiveSupport::ExecutionContext[:temporary] = "value"
     ActiveRecord::QueryLogs.tags = [ temporary_tag: ->(context) { context[:temporary] } ]
 
     assert_equal "/*temporary_tag:value*/", ActiveRecord::QueryLogs.call("")
 
-    ActiveRecord::QueryLogs.set_context(temporary: "new_value")
+    ActiveSupport::ExecutionContext[:temporary] = "new_value"
 
     assert_nil ActiveRecord::QueryLogs.cached_comment
     assert_equal "/*temporary_tag:new_value*/", ActiveRecord::QueryLogs.call("")
-  ensure
-    ActiveRecord::QueryLogs.cached_comment = nil
-    ActiveRecord::QueryLogs.cache_query_log_tags = false
-  end
-
-  def test_ensure_context_has_symbol_keys
-    ActiveRecord::QueryLogs.tags = [ new_key: ->(context) { context[:symbol_key] } ]
-    ActiveRecord::QueryLogs.set_context("symbol_key" => "symbolized")
-
-    assert_sql(%r{/\*new_key:symbolized}) do
-      Dashboard.first
-    end
-  ensure
-    ActiveRecord::QueryLogs.set_context(application_name: nil)
   end
 
   def test_default_tag_behavior
     ActiveRecord::QueryLogs.tags = [:application, :foo]
-    ActiveRecord::QueryLogs.set_context(foo: "bar") do
+    ActiveSupport::ExecutionContext.set(foo: "bar") do
       assert_sql(%r{/\*application:active_record,foo:bar\*/}) do
         Dashboard.first
       end
@@ -157,39 +142,29 @@ class QueryLogsTest < ActiveRecord::TestCase
   end
 
   def test_empty_comments_are_not_added
-    original_tags = ActiveRecord::QueryLogs.tags
     ActiveRecord::QueryLogs.tags = [ empty: -> { nil } ]
     assert_sql(%r{select id from posts$}) do
       ActiveRecord::Base.connection.execute "select id from posts"
     end
-  ensure
-    ActiveRecord::QueryLogs.tags = original_tags
   end
 
   def test_custom_basic_tags
-    original_tags = ActiveRecord::QueryLogs.tags
     ActiveRecord::QueryLogs.tags = [ :application, { custom_string: "test content" } ]
 
     assert_sql(%r{/\*application:active_record,custom_string:test content\*/$}) do
       Dashboard.first
     end
-  ensure
-    ActiveRecord::QueryLogs.tags = original_tags
   end
 
   def test_custom_proc_tags
-    original_tags = ActiveRecord::QueryLogs.tags
     ActiveRecord::QueryLogs.tags = [ :application, { custom_proc: -> { "test content" } } ]
 
     assert_sql(%r{/\*application:active_record,custom_proc:test content\*/$}) do
       Dashboard.first
     end
-  ensure
-    ActiveRecord::QueryLogs.tags = original_tags
   end
 
   def test_multiple_custom_tags
-    original_tags = ActiveRecord::QueryLogs.tags
     ActiveRecord::QueryLogs.tags = [
       :application,
       { custom_proc: -> { "test content" }, another_proc: -> { "more test content" } },
@@ -198,34 +173,14 @@ class QueryLogsTest < ActiveRecord::TestCase
     assert_sql(%r{/\*application:active_record,custom_proc:test content,another_proc:more test content\*/$}) do
       Dashboard.first
     end
-  ensure
-    ActiveRecord::QueryLogs.tags = original_tags
   end
 
   def test_custom_proc_context_tags
-    original_tags = ActiveRecord::QueryLogs.tags
-    ActiveRecord::QueryLogs.set_context(foo: "bar")
+    ActiveSupport::ExecutionContext[:foo] = "bar"
     ActiveRecord::QueryLogs.tags = [ :application, { custom_context_proc: ->(context) { context[:foo] } } ]
 
     assert_sql(%r{/\*application:active_record,custom_context_proc:bar\*/$}) do
       Dashboard.first
     end
-  ensure
-    ActiveRecord::QueryLogs.set_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) { 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
diff --git a/activesupport/lib/active_support.rb b/activesupport/lib/active_support.rb
index 3e35864cc5..f148992690 100644
--- a/activesupport/lib/active_support.rb
+++ b/activesupport/lib/active_support.rb
@@ -40,6 +40,7 @@ module ActiveSupport
   autoload :CurrentAttributes
   autoload :Dependencies
   autoload :DescendantsTracker
+  autoload :ExecutionContext
   autoload :ExecutionWrapper
   autoload :Executor
   autoload :FileUpdateChecker
diff --git a/activesupport/lib/active_support/execution_context.rb b/activesupport/lib/active_support/execution_context.rb
new file mode 100644
index 0000000000..01d7742e9a
--- /dev/null
+++ b/activesupport/lib/active_support/execution_context.rb
@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+module ActiveSupport
+  module ExecutionContext # :nodoc:
+    @after_change_callbacks = []
+    class << self
+      def after_change(&block)
+        @after_change_callbacks << block
+      end
+
+      # Updates the execution context. If a block is given, it resets the provided keys to their
+      # previous value once the block exits.
+      def set(**options)
+        options.symbolize_keys!
+        keys = options.keys
+
+        store = self.store
+
+        previous_context = keys.zip(store.values_at(*keys)).to_h
+
+        store.merge!(options)
+        @after_change_callbacks.each(&:call)
+
+        if block_given?
+          begin
+            yield
+          ensure
+            store.merge!(previous_context)
+            @after_change_callbacks.each(&:call)
+          end
+        end
+      end
+
+      def []=(key, value)
+        store[key.to_sym] = value
+        @after_change_callbacks.each(&:call)
+      end
+
+      def to_h
+        store.dup
+      end
+
+      def clear
+        store.clear
+      end
+
+      private
+        def store
+          Thread.current[:active_support_execution_context] ||= {}
+        end
+    end
+  end
+end
diff --git a/activesupport/lib/active_support/execution_context/test_helper.rb b/activesupport/lib/active_support/execution_context/test_helper.rb
new file mode 100644
index 0000000000..ae8c43ea0c
--- /dev/null
+++ b/activesupport/lib/active_support/execution_context/test_helper.rb
@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module ActiveSupport::ExecutionContext::TestHelper # :nodoc:
+  def before_setup
+    ActiveSupport::ExecutionContext.clear
+    super
+  end
+
+  def after_teardown
+    super
+    ActiveSupport::ExecutionContext.clear
+  end
+end
diff --git a/activesupport/lib/active_support/railtie.rb b/activesupport/lib/active_support/railtie.rb
index e4892e0308..30177ac472 100644
--- a/activesupport/lib/active_support/railtie.rb
+++ b/activesupport/lib/active_support/railtie.rb
@@ -27,6 +27,12 @@ module ActiveSupport
       end
     end
 
+    initializer "active_support.reset_execution_context" do |app|
+      app.reloader.before_class_unload { ActiveSupport::ExecutionContext.clear }
+      app.executor.to_run              { ActiveSupport::ExecutionContext.clear }
+      app.executor.to_complete         { ActiveSupport::ExecutionContext.clear }
+    end
+
     initializer "active_support.reset_all_current_attributes_instances" do |app|
       executor_around_test_case = app.config.active_support.delete(:executor_around_test_case)
 
@@ -41,6 +47,9 @@ module ActiveSupport
         else
           require "active_support/current_attributes/test_helper"
           include ActiveSupport::CurrentAttributes::TestHelper
+
+          require "active_support/execution_context/test_helper"
+          include ActiveSupport::ExecutionContext::TestHelper
         end
       end
     end
diff --git a/activesupport/test/current_attributes_test.rb b/activesupport/test/current_attributes_test.rb
index 088c041d5e..85cb111dee 100644
--- a/activesupport/test/current_attributes_test.rb
+++ b/activesupport/test/current_attributes_test.rb
@@ -1,17 +1,12 @@
 # frozen_string_literal: true
 
 require_relative "abstract_unit"
+require "active_support/current_attributes/test_helper"
 
 class CurrentAttributesTest < ActiveSupport::TestCase
   # CurrentAttributes is automatically reset in Rails app via executor hooks set in railtie
   # But not in Active Support's own test suite.
-  setup do
-    ActiveSupport::CurrentAttributes.reset_all
-  end
-
-  teardown do
-    ActiveSupport::CurrentAttributes.reset_all
-  end
+  include ActiveSupport::CurrentAttributes::TestHelper
 
   Person = Struct.new(:id, :name, :time_zone)
 
diff --git a/activesupport/test/execution_context_test.rb b/activesupport/test/execution_context_test.rb
new file mode 100644
index 0000000000..502044e4b2
--- /dev/null
+++ b/activesupport/test/execution_context_test.rb
@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+require_relative "abstract_unit"
+require "active_support/execution_context/test_helper"
+
+class ExecutionContextTest < ActiveSupport::TestCase
+  # ExecutionContext is automatically reset in Rails app via executor hooks set in railtie
+  # But not in Active Support's own test suite.
+  include ActiveSupport::ExecutionContext::TestHelper
+
+  test "#set restore the modified keys when the block exits" do
+    assert_nil ActiveSupport::ExecutionContext.to_h[:foo]
+    ActiveSupport::ExecutionContext.set(foo: "bar") do
+      assert_equal "bar", ActiveSupport::ExecutionContext.to_h[:foo]
+      ActiveSupport::ExecutionContext.set(foo: "plop") do
+        assert_equal "plop", ActiveSupport::ExecutionContext.to_h[:foo]
+      end
+      assert_equal "bar", ActiveSupport::ExecutionContext.to_h[:foo]
+
+      ActiveSupport::ExecutionContext[:direct_assignment] = "present"
+      ActiveSupport::ExecutionContext.set(multi_assignment: "present")
+    end
+
+    assert_nil ActiveSupport::ExecutionContext.to_h[:foo]
+
+    assert_equal "present", ActiveSupport::ExecutionContext.to_h[:direct_assignment]
+    assert_equal "present", ActiveSupport::ExecutionContext.to_h[:multi_assignment]
+  end
+
+  test "#set coerce keys to symbol" do
+    ActiveSupport::ExecutionContext.set("foo" => "bar") do
+      assert_equal "bar", ActiveSupport::ExecutionContext.to_h[:foo]
+    end
+  end
+
+  test "#[]= coerce keys to symbol" do
+    ActiveSupport::ExecutionContext["symbol_key"] = "symbolized"
+    assert_equal "symbolized", ActiveSupport::ExecutionContext.to_h[:symbol_key]
+  end
+
+  test "#to_h returns a copy of the context" do
+    ActiveSupport::ExecutionContext[:foo] = 42
+    context = ActiveSupport::ExecutionContext.to_h
+    context[:foo] = 43
+    assert_equal 42, ActiveSupport::ExecutionContext.to_h[:foo]
+  end
+end