Merge pull request rails#43079 from Shopify/ar-query-logs-instance-exec

Pass `ActiveRecord::QueryLogs` as argument rather than `instance_exec`

Author: byroot

actionpack/lib/action_controller/railtie.rb

@@ -115,9 +115,9 @@ class Railtie < Rails::Railtie # :nodoc:
 
         ActiveSupport.on_load(:active_record) do
           ActiveRecord::QueryLogs.taggings.merge!(
-            controller:            -> { context[:controller].controller_name },
-            action:                -> { context[:controller].action_name },
-            namespaced_controller: -> { context[:controller].class.name }
+            controller:            ->(context) { context[:controller].controller_name },
+            action:                ->(context) { context[:controller].action_name },
+            namespaced_controller: ->(context) { context[:controller].class.name }
           )
         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[:job].class.name }
+          ActiveRecord::QueryLogs.taggings[:job] = ->(context) { context[:job].class.name }
         end
       end
     end

activerecord/lib/active_record/query_logs.rb

@@ -37,7 +37,10 @@ module ActiveRecord
   #
   #    tags = [
   #      :application,
-  #      { custom_tag: -> { context[:controller].controller_name } }
+  #      {
+  #        custom_tag: ->(context) { context[:controller].controller_name },
+  #        custom_value: -> { Custom.value },
+  #      }
   #    ]
   #    ActiveRecord::QueryLogs.tags = tags
   #
@@ -173,30 +176,24 @@ def escape_sql_comment(content)
 
         def tag_content
           tags.flat_map { |i| [*i] }.filter_map do |tag|
-            key, value_input = tag
-
-            val = case value_input
-                  when nil then tag_value(key)
-                  when Proc then instance_exec(&value_input)
-                  else value_input
+            key, handler = tag
+            handler ||= taggings[key]
+
+            val = if handler.nil?
+              context[key]
+            elsif handler.respond_to?(:call)
+              if handler.arity == 0
+                handler.call
+              else
+                handler.call(context)
+              end
+            else
+              handler
             end
-
             "#{key}:#{val}" unless val.nil?
           end.join(",")
         end
 
-        def tag_value(key)
-          if value = taggings[key]
-            if value.respond_to?(:call)
-              instance_exec(&taggings[key])
-            else
-              value
-            end
-          else
-            context[key]
-          end
-        end
-
         def inline_tag_content
           inline_tags.join
         end

activerecord/test/cases/query_logs_test.rb

@@ -113,7 +113,7 @@ 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.tags = [ temporary_tag: -> { context[:temporary] } ]
+    ActiveRecord::QueryLogs.tags = [ temporary_tag: ->(context) { context[:temporary] } ]
 
     assert_equal " /*temporary_tag:value*/", ActiveRecord::QueryLogs.add_query_log_tags_to_sql("")
 
@@ -127,7 +127,7 @@ def test_resets_cache_on_context_update
   end
 
   def test_ensure_context_has_symbol_keys
-    ActiveRecord::QueryLogs.tags = [ new_key: -> { context[:symbol_key] } ]
+    ActiveRecord::QueryLogs.tags = [ new_key: ->(context) { context[:symbol_key] } ]
     ActiveRecord::QueryLogs.update_context("symbol_key" => "symbolized")
 
     assert_sql(%r{/\*new_key:symbolized}) do
@@ -244,7 +244,10 @@ def test_custom_proc_tags
 
   def test_multiple_custom_tags
     original_tags = ActiveRecord::QueryLogs.tags
-    ActiveRecord::QueryLogs.tags = [ :application, { custom_proc: -> { "test content" }, another_proc: -> { "more test content" } } ]
+    ActiveRecord::QueryLogs.tags = [
+      :application,
+      { custom_proc: -> { "test content" }, another_proc: -> { "more test content" } },
+    ]
 
     assert_sql(%r{/\*application:active_record,custom_proc:test content,another_proc:more test content\*/$}) do
       Dashboard.first
@@ -256,7 +259,7 @@ def test_multiple_custom_tags
   def test_custom_proc_context_tags
     original_tags = ActiveRecord::QueryLogs.tags
     ActiveRecord::QueryLogs.update_context(foo: "bar")
-    ActiveRecord::QueryLogs.tags = [ :application, { custom_context_proc: -> { context[:foo] } } ]
+    ActiveRecord::QueryLogs.tags = [ :application, { custom_context_proc: ->(context) { context[:foo] } } ]
 
     assert_sql(%r{/\*application:active_record,custom_context_proc:bar\*/$}) do
       Dashboard.first
@@ -268,7 +271,7 @@ def test_custom_proc_context_tags
 
   def test_set_context_restore_state
     original_tags = ActiveRecord::QueryLogs.tags
-    ActiveRecord::QueryLogs.tags = [foo: -> { context[:foo] }]
+    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

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[: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[:job].dynamic_content } } ]"
+      add_to_config "config.active_record.query_log_tags = [ { dynamic: ->(context) { context[:job].dynamic_content } } ]"
 
       boot_app