From 6bad959565900499325d4239266cbb1f1f8e4056 Mon Sep 17 00:00:00 2001 From: Jean Boussier 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 61fc35561eee9..11403d8f7dddc 100644 --- a/actionpack/lib/abstract_controller/base.rb +++ b/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 diff --git a/actionpack/lib/action_controller.rb b/actionpack/lib/action_controller.rb index d80b4f5b030a7..62ae8f8d9eed3 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 f8ca992f0f250..0000000000000 --- 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 bea6873a503c9..701ddb5befab5 100644 --- a/actionpack/lib/action_controller/railtie.rb +++ b/actionpack/lib/action_controller/railtie.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 }, diff --git a/activejob/lib/active_job/execution.rb b/activejob/lib/active_job/execution.rb index 6f07c021221d8..12c1b468c3c25 100644 --- a/activejob/lib/active_job/execution.rb +++ b/activejob/lib/active_job/execution.rb @@ -54,6 +54,7 @@ def perform(*) 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 08b929dca9f27..0000000000000 --- 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 7755b66e285d3..27fb31062c427 100644 --- a/activejob/lib/active_job/railtie.rb +++ b/activejob/lib/active_job/railtie.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 diff --git a/activerecord/lib/active_record/query_logs.rb b/activerecord/lib/active_record/query_logs.rb index 452cdb00879e7..f116a154dd3cc 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 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 cache_query_log_tags 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] diff --git a/activerecord/lib/active_record/railtie.rb b/activerecord/lib/active_record/railtie.rb index 950ac63eb6611..04112bc2e170b 100644 --- a/activerecord/lib/active_record/railtie.rb +++ b/activerecord/lib/active_record/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 diff --git a/activerecord/test/cases/query_logs_test.rb b/activerecord/test/cases/query_logs_test.rb index cc34f3a39b1c7..05207f17e3d2f 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 @@ 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 diff --git a/activesupport/lib/active_support.rb b/activesupport/lib/active_support.rb index 3e35864cc5900..f148992690f62 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 0000000000000..01d7742e9ad62 --- /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 0000000000000..ae8c43ea0c30e --- /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 e4892e03083f5..30177ac472cca 100644 --- a/activesupport/lib/active_support/railtie.rb +++ b/activesupport/lib/active_support/railtie.rb @@ -27,6 +27,12 @@ class Railtie < Rails::Railtie # :nodoc: 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 @@ class Railtie < Rails::Railtie # :nodoc: 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 088c041d5e0dd..85cb111deeef1 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 0000000000000..502044e4b20f2 --- /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