Merge pull request rails#43535 from Shopify/remove-query-logs-null-object

Refactor ActiveRecord::QueryLogs context API

Author: byroot

actionpack/lib/action_controller/railtie.rb

@@ -119,9 +119,9 @@ class Railtie < Rails::Railtie # :nodoc:
 
         ActiveSupport.on_load(:active_record) do
           ActiveRecord::QueryLogs.taggings.merge!(
-            controller:            ->(context) { context[:controller].controller_name },
-            action:                ->(context) { context[:controller].action_name },
-            namespaced_controller: ->(context) { context[:controller].class.name }
+            controller:            ->(context) { context[:controller]&.controller_name },
+            action:                ->(context) { context[:controller]&.action_name },
+            namespaced_controller: ->(context) { context[:controller].class.name if context[:controller] }
           )
         end
       end

activejob/lib/active_job/railtie.rb

@@ -70,7 +70,7 @@ class Railtie < Rails::Railtie # :nodoc:
         end
 
         ActiveSupport.on_load(:active_record) do
-          ActiveRecord::QueryLogs.taggings[:job] = ->(context) { context[:job].class.name unless context[:job].nil? }
+          ActiveRecord::QueryLogs.taggings[:job] = ->(context) { context[:job].class.name if context[:job] }
         end
       end
     end

activerecord/lib/active_record/query_logs.rb

@@ -38,24 +38,24 @@ module ActiveRecord
   #    tags = [
   #      :application,
   #      {
-  #        custom_tag: ->(context) { context[:controller].controller_name },
+  #        custom_tag: ->(context) { context[:controller]&.controller_name },
   #        custom_value: -> { Custom.value },
   #      }
   #    ]
   #    ActiveRecord::QueryLogs.tags = tags
   #
-  # The QueryLogs +context+ can be manipulated via +update_context+ & +set_context+ methods.
-  #
-  # Direct updates to a context value:
-  #
-  #    ActiveRecord::QueryLogs.update_context(foo: Bar.new)
+  # The QueryLogs +context+ can be manipulated via the +set_context+ method.
   #
   # Temporary updates limited to the execution of a block:
   #
   #    ActiveRecord::QueryLogs.set_context(foo: Bar.new) do
   #      posts = Post.all
   #    end
   #
+  # Direct updates to a context value:
+  #
+  #    ActiveRecord::QueryLogs.set_context(foo: Bar.new)
+  #
   # Tag comments can be prepended to the query:
   #
   #    ActiveRecord::QueryLogs.prepend_comment = true
@@ -75,39 +75,25 @@ module QueryLogs
     mattr_accessor :cache_query_log_tags, instance_accessor: false, default: false
     thread_mattr_accessor :cached_comment, instance_accessor: false
 
-    class NullObject # :nodoc:
-      def method_missing(method, *args, &block)
-        NullObject.new
-      end
-
-      def nil?
-        true
-      end
-
-      private
-        def respond_to_missing?(method, include_private = false)
-          true
-        end
-    end
-
     class << self
       # Updates the context used to construct tags in the SQL comment.
-      # Resets the cached comment if <tt>cache_query_log_tags</tt> is +true+.
-      def update_context(**options)
-        context.merge!(**options.symbolize_keys)
-        self.cached_comment = nil
-      end
-
-      # Updates the context used to construct tags in the SQL comment during
-      # execution of the provided block. Resets the provided keys to their
+      # 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
-        update_context(**options)
-        yield if block_given?
-      ensure
-        update_context(**previous_context)
+        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 call(sql) # :nodoc:
@@ -137,7 +123,7 @@ def uncached_comment
         end
 
         def context
-          Thread.current[:active_record_query_log_tags_context] ||= Hash.new { NullObject.new }
+          Thread.current[:active_record_query_log_tags_context] ||= {}
         end
 
         def escape_sql_comment(content)

activerecord/test/cases/query_logs_test.rb

@@ -112,12 +112,12 @@ def test_retrieves_comment_from_cache_when_enabled_and_set
 
   def test_resets_cache_on_context_update
     ActiveRecord::QueryLogs.cache_query_log_tags = true
-    ActiveRecord::QueryLogs.update_context(temporary: "value")
+    ActiveRecord::QueryLogs.set_context(temporary: "value")
     ActiveRecord::QueryLogs.tags = [ temporary_tag: ->(context) { context[:temporary] } ]
 
     assert_equal "/*temporary_tag:value*/", ActiveRecord::QueryLogs.call("")
 
-    ActiveRecord::QueryLogs.update_context(temporary: "new_value")
+    ActiveRecord::QueryLogs.set_context(temporary: "new_value")
 
     assert_nil ActiveRecord::QueryLogs.cached_comment
     assert_equal "/*temporary_tag:new_value*/", ActiveRecord::QueryLogs.call("")
@@ -128,13 +128,13 @@ def test_resets_cache_on_context_update
 
   def test_ensure_context_has_symbol_keys
     ActiveRecord::QueryLogs.tags = [ new_key: ->(context) { context[:symbol_key] } ]
-    ActiveRecord::QueryLogs.update_context("symbol_key" => "symbolized")
+    ActiveRecord::QueryLogs.set_context("symbol_key" => "symbolized")
 
     assert_sql(%r{/\*new_key:symbolized}) do
       Dashboard.first
     end
   ensure
-    ActiveRecord::QueryLogs.update_context(application_name: nil)
+    ActiveRecord::QueryLogs.set_context(application_name: nil)
   end
 
   def test_default_tag_behavior
@@ -197,14 +197,14 @@ def test_multiple_custom_tags
 
   def test_custom_proc_context_tags
     original_tags = ActiveRecord::QueryLogs.tags
-    ActiveRecord::QueryLogs.update_context(foo: "bar")
+    ActiveRecord::QueryLogs.set_context(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.update_context(foo: nil)
+    ActiveRecord::QueryLogs.set_context(foo: nil)
     ActiveRecord::QueryLogs.tags = original_tags
   end
 

railties/test/application/query_logs_test.rb

@@ -123,7 +123,7 @@ def app
     test "query cache is cleared between requests" do
       add_to_config "config.active_record.query_log_tags_enabled = true"
       add_to_config "config.active_record.cache_query_log_tags = true"
-      add_to_config "config.active_record.query_log_tags = [ { dynamic: ->(context) { context[:controller].dynamic_content } } ]"
+      add_to_config "config.active_record.query_log_tags = [ { dynamic: ->(context) { context[:controller]&.dynamic_content } } ]"
 
       boot_app
 
@@ -141,7 +141,7 @@ def app
     test "query cache is cleared between job executions" do
       add_to_config "config.active_record.query_log_tags_enabled = true"
       add_to_config "config.active_record.cache_query_log_tags = true"
-      add_to_config "config.active_record.query_log_tags = [ { dynamic: ->(context) { context[:job].dynamic_content } } ]"
+      add_to_config "config.active_record.query_log_tags = [ { dynamic: ->(context) { context[:job]&.dynamic_content } } ]"
 
       boot_app