Merge pull request rails#52392 from Shopify/perf-query-logs

Optimize ActiveRecord::QueryLogs

Author: byroot

activerecord/lib/active_record/query_logs.rb

@@ -72,14 +72,70 @@ module ActiveRecord
   #
   #    config.active_record.cache_query_log_tags = true
   module QueryLogs
-    mattr_accessor :taggings, instance_accessor: false, default: {}
-    mattr_accessor :tags, instance_accessor: false, default: [ :application ]
-    mattr_accessor :prepend_comment, instance_accessor: false, default: false
-    mattr_accessor :cache_query_log_tags, instance_accessor: false, default: false
-    mattr_accessor :tags_formatter, instance_accessor: false
+    class GetKeyHandler # :nodoc:
+      def initialize(name)
+        @name = name
+      end
+
+      def call(context)
+        context[@name]
+      end
+    end
+
+    class IdentityHandler # :nodoc:
+      def initialize(value)
+        @value = value
+      end
+
+      def call(_context)
+        @value
+      end
+    end
+
+    class ZeroArityHandler # :nodoc:
+      def initialize(proc)
+        @proc = proc
+      end
+
+      def call(_context)
+        @proc.call
+      end
+    end
+
+    @taggings = {}
+    @tags = [ :application ]
+    @prepend_comment = false
+    @cache_query_log_tags = false
+    @tags_formatter = false
+
     thread_mattr_accessor :cached_comment, instance_accessor: false
 
     class << self
+      attr_reader :tags, :taggings, :tags_formatter # :nodoc:
+      attr_accessor :prepend_comment, :cache_query_log_tags # :nodoc:
+
+      def taggings=(taggings) # :nodoc:
+        @taggings = taggings
+        @handlers = rebuild_handlers
+      end
+
+      def tags=(tags) # :nodoc:
+        @tags = tags
+        @handlers = rebuild_handlers
+      end
+
+      def tags_formatter=(format) # :nodoc:
+        @tags_formatter = format
+        @formatter = case format
+        when :legacy
+          LegacyFormatter
+        when :sqlcommenter
+          SQLCommenter
+        else
+          raise ArgumentError, "Formatter is unsupported: #{format}"
+        end
+      end
+
       def call(sql, connection) # :nodoc:
         comment = self.comment(connection)
 
@@ -96,19 +152,6 @@ def clear_cache # :nodoc:
         self.cached_comment = nil
       end
 
-      # Updates the formatter to be what the passed in format is.
-      def update_formatter(format)
-        self.tags_formatter =
-          case format
-          when :legacy
-            LegacyFormatter.new
-          when :sqlcommenter
-            SQLCommenter.new
-          else
-            raise ArgumentError, "Formatter is unsupported: #{formatter}"
-          end
-      end
-
       if Thread.respond_to?(:each_caller_location)
         def query_source_location # :nodoc:
           Thread.each_caller_location do |location|
@@ -126,6 +169,36 @@ def query_source_location # :nodoc:
       ActiveSupport::ExecutionContext.after_change { ActiveRecord::QueryLogs.clear_cache }
 
       private
+        def rebuild_handlers
+          handlers = []
+          @tags.each do |i|
+            if i.is_a?(Hash)
+              i.each do |k, v|
+                handlers << [k, build_handler(k, v)]
+              end
+            else
+              handlers << [i, build_handler(i)]
+            end
+          end
+          handlers.sort_by! { |(key, _)| key.to_s }
+          handlers
+        end
+
+        def build_handler(name, handler = nil)
+          handler ||= @taggings[name]
+          if handler.nil?
+            GetKeyHandler.new(name)
+          elsif handler.respond_to?(:call)
+            if handler.arity == 0
+              ZeroArityHandler.new(handler)
+            else
+              handler
+            end
+          else
+            IdentityHandler.new(handler)
+          end
+        end
+
         # 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+.
         def comment(connection)
@@ -136,10 +209,6 @@ def comment(connection)
           end
         end
 
-        def formatter
-          self.tags_formatter || self.update_formatter(:legacy)
-        end
-
         def uncached_comment(connection)
           content = tag_content(connection)
 
@@ -165,25 +234,15 @@ def tag_content(connection)
           context = ActiveSupport::ExecutionContext.to_h
           context[:connection] ||= connection
 
-          pairs = tags.flat_map { |i| [*i] }.filter_map do |tag|
-            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?
+          pairs = @handlers.filter_map do |(key, handler)|
+            val = handler.call(context)
+            @formatter.format(key, val) unless val.nil?
           end
-          self.formatter.format(pairs)
+          @formatter.join(pairs)
         end
     end
+
+    @handlers = rebuild_handlers
+    self.tags_formatter = :legacy
   end
 end

activerecord/lib/active_record/query_logs_formatter.rb

@@ -2,40 +2,29 @@
 
 module ActiveRecord
   module QueryLogs
-    class LegacyFormatter # :nodoc:
-      def initialize
-        @key_value_separator = ":"
-      end
-
-      # Formats the key value pairs into a string.
-      def format(pairs)
-        pairs.map! do |key, value|
-          "#{key}#{key_value_separator}#{format_value(value)}"
-        end.join(",")
-      end
-
-      private
-        attr_reader :key_value_separator
-
-        def format_value(value)
-          value
+    module LegacyFormatter # :nodoc:
+      class << self
+        # Formats the key value pairs into a string.
+        def format(key, value)
+          "#{key}:#{value}"
         end
-    end
 
-    class SQLCommenter < LegacyFormatter # :nodoc:
-      def initialize
-        @key_value_separator = "="
+        def join(pairs)
+          pairs.join(",")
+        end
       end
+    end
 
-      def format(pairs)
-        pairs.sort_by! { |pair| pair.first.to_s }
-        super
-      end
+    class SQLCommenter # :nodoc:
+      class << self
+        def format(key, value)
+          "#{key}='#{ERB::Util.url_encode(value)}'"
+        end
 
-      private
-        def format_value(value)
-          "'#{ERB::Util.url_encode(value)}'"
+        def join(pairs)
+          pairs.join(",")
         end
+      end
     end
   end
 end

activerecord/test/cases/query_logs_test.rb

@@ -33,7 +33,7 @@ def teardown
     ActiveRecord::QueryLogs.prepend_comment = false
     ActiveRecord::QueryLogs.cache_query_log_tags = false
     ActiveRecord::QueryLogs.clear_cache
-    ActiveRecord::QueryLogs.update_formatter(:legacy)
+    ActiveRecord::QueryLogs.tags_formatter = :legacy
 
     # 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.
@@ -45,7 +45,8 @@ def test_escaping_good_comment
   end
 
   def test_escaping_good_comment_with_custom_separator
-    ActiveRecord::QueryLogs.update_formatter(:sqlcommenter)
+    ActiveRecord::QueryLogs.tags_formatter = :sqlcommenter
+
     assert_equal "app='foo'", ActiveRecord::QueryLogs.send(:escape_sql_comment, "app='foo'")
   end
 
@@ -183,7 +184,8 @@ def test_empty_comments_are_not_added
   end
 
   def test_sql_commenter_format
-    ActiveRecord::QueryLogs.update_formatter(:sqlcommenter)
+    ActiveRecord::QueryLogs.tags_formatter = :sqlcommenter
+
     assert_queries_match(%r{/\*application='active_record'\*/}) do
       Dashboard.first
     end
@@ -211,13 +213,13 @@ def test_multiple_custom_tags
       { custom_proc: -> { "test content" }, another_proc: -> { "more test content" } },
     ]
 
-    assert_queries_match(%r{/\*application:active_record,custom_proc:test content,another_proc:more test content\*/}) do
+    assert_queries_match(%r{/\*another_proc:more test content,application:active_record,custom_proc:test content\*/}) do
       Dashboard.first
     end
   end
 
   def test_sqlcommenter_format_value
-    ActiveRecord::QueryLogs.update_formatter(:sqlcommenter)
+    ActiveRecord::QueryLogs.tags_formatter = :sqlcommenter
 
     ActiveRecord::QueryLogs.tags = [
       :application,
@@ -230,7 +232,7 @@ def test_sqlcommenter_format_value
   end
 
   def test_sqlcommenter_format_allows_string_keys
-    ActiveRecord::QueryLogs.update_formatter(:sqlcommenter)
+    ActiveRecord::QueryLogs.tags_formatter = :sqlcommenter
 
     ActiveRecord::QueryLogs.tags = [
       :application,
@@ -247,7 +249,7 @@ def test_sqlcommenter_format_allows_string_keys
   end
 
   def test_sqlcommenter_format_value_string_coercible
-    ActiveRecord::QueryLogs.update_formatter(:sqlcommenter)
+    ActiveRecord::QueryLogs.tags_formatter = :sqlcommenter
 
     ActiveRecord::QueryLogs.tags = [
       :application,