Skip to content

Commit

Permalink
Merge pull request rails#52392 from Shopify/perf-query-logs
Browse files Browse the repository at this point in the history
Optimize ActiveRecord::QueryLogs
  • Loading branch information
byroot authored Jul 22, 2024
2 parents 2fdf772 + 6fb612b commit b291408
Show file tree
Hide file tree
Showing 3 changed files with 124 additions and 74 deletions.
137 changes: 98 additions & 39 deletions activerecord/lib/active_record/query_logs.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand All @@ -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|
Expand All @@ -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)
Expand All @@ -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)

Expand All @@ -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
45 changes: 17 additions & 28 deletions activerecord/lib/active_record/query_logs_formatter.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
16 changes: 9 additions & 7 deletions activerecord/test/cases/query_logs_test.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down

0 comments on commit b291408

Please sign in to comment.