Merge pull request rails#43598 from Shopify/active-support-execution-context

Extract ActiveSupport::ExecutionContext out of ActiveRecord::QueryLogs

Author: byroot

actionpack/lib/abstract_controller/base.rb

@@ -148,6 +148,7 @@ def process(action, *args)
 
       @_response_body = nil
 
+      ActiveSupport::ExecutionContext[:controller] = self
       process_action(action_name, *args)
     end
 

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

actionpack/lib/action_controller/metal/query_tags.rb

@@ -113,10 +113,6 @@ class Railtie < Rails::Railtie # :nodoc:
       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 },

actionpack/lib/action_controller/railtie.rb

@@ -54,6 +54,7 @@ def perform(*)
 
     private
       def _perform_job
+        ActiveSupport::ExecutionContext[:job] = self
         run_callbacks :perform do
           perform(*arguments)
         end

activejob/lib/active_job/execution.rb

@@ -65,10 +65,6 @@ class Railtie < Rails::Railtie # :nodoc:
       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

activejob/lib/active_job/query_tags.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 QueryLogs
     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 @@ def call(sql) # :nodoc:
         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 @@ def uncached_comment
           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]

activejob/lib/active_job/railtie.rb

@@ -376,10 +376,6 @@ class Railtie < Rails::Railtie # :nodoc:
           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

activerecord/lib/active_record/query_logs.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 @@ def setup
   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 @@ def test_add_comments_to_beginning_of_query
     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 @@ def test_retrieves_comment_from_cache_when_enabled_and_set
     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 @@ def test_default_tag_behavior
   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 @@ def test_multiple_custom_tags
     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