|
| 1 | +# frozen_string_literal: true |
| 2 | + |
| 3 | +require "active_support/core_ext/module/attribute_accessors_per_thread" |
| 4 | + |
| 5 | +module ActiveRecord |
| 6 | + # = Active Record Query Logs |
| 7 | + # |
| 8 | + # Automatically tag SQL queries with runtime information. |
| 9 | + # |
| 10 | + # Default tags available for use: |
| 11 | + # |
| 12 | + # * +application+ |
| 13 | + # * +pid+ |
| 14 | + # * +socket+ |
| 15 | + # * +db_host+ |
| 16 | + # * +database+ |
| 17 | + # |
| 18 | + # _Action Controller and Active Job tags are also defined when used in Rails:_ |
| 19 | + # |
| 20 | + # * +controller+ |
| 21 | + # * +action+ |
| 22 | + # * +job+ |
| 23 | + # |
| 24 | + # The tags used in a query can be configured directly: |
| 25 | + # |
| 26 | + # ActiveRecord::QueryLogs.tags = [ :application, :controller, :action, :job ] |
| 27 | + # |
| 28 | + # or via Rails configuration: |
| 29 | + # |
| 30 | + # config.active_record.query_log_tags = [ :application, :controller, :action, :job ] |
| 31 | + # |
| 32 | + # To add new comment tags, add a hash to the tags array containing the keys and values you |
| 33 | + # want to add to the comment. Dynamic content can be created by setting a proc or lambda value in a hash, |
| 34 | + # and can reference any value stored in the +context+ object. |
| 35 | + # |
| 36 | + # Example: |
| 37 | + # |
| 38 | + # tags = [ |
| 39 | + # :application, |
| 40 | + # { custom_tag: -> { context[:controller].controller_name } } |
| 41 | + # ] |
| 42 | + # ActiveRecord::QueryLogs.tags = tags |
| 43 | + # |
| 44 | + # The QueryLogs +context+ can be manipulated via +update_context+ & +set_context+ methods. |
| 45 | + # |
| 46 | + # Direct updates to a context value: |
| 47 | + # |
| 48 | + # ActiveRecord::QueryLogs.update_context(foo: Bar.new) |
| 49 | + # |
| 50 | + # Temporary updates limited to the execution of a block: |
| 51 | + # |
| 52 | + # ActiveRecord::QueryLogs.set_context(foo: Bar.new) do |
| 53 | + # posts = Post.all |
| 54 | + # end |
| 55 | + # |
| 56 | + # Tag comments can be prepended to the query: |
| 57 | + # |
| 58 | + # ActiveRecord::QueryLogs.prepend_comment = true |
| 59 | + # |
| 60 | + # For applications where the content will not change during the lifetime of |
| 61 | + # the request or job execution, the tags can be cached for reuse in every query: |
| 62 | + # |
| 63 | + # ActiveRecord::QueryLogs.cache_query_log_tags = true |
| 64 | + # |
| 65 | + # This option can be set during application configuration or in a Rails initializer: |
| 66 | + # |
| 67 | + # config.active_record.cache_query_log_tags = true |
| 68 | + module QueryLogs |
| 69 | + mattr_accessor :taggings, instance_accessor: false, default: {} |
| 70 | + mattr_accessor :tags, instance_accessor: false, default: [ :application ] |
| 71 | + mattr_accessor :prepend_comment, instance_accessor: false, default: false |
| 72 | + mattr_accessor :cache_query_log_tags, instance_accessor: false, default: false |
| 73 | + thread_mattr_accessor :cached_comment, instance_accessor: false |
| 74 | + |
| 75 | + class << self |
| 76 | + # Updates the context used to construct tags in the SQL comment. |
| 77 | + # Resets the cached comment if <tt>cache_query_log_tags</tt> is +true+. |
| 78 | + def update_context(**options) |
| 79 | + context.merge!(**options.symbolize_keys) |
| 80 | + self.cached_comment = nil |
| 81 | + end |
| 82 | + |
| 83 | + # Updates the context used to construct tags in the SQL comment during |
| 84 | + # execution of the provided block. Resets provided values to nil after |
| 85 | + # the block is executed. |
| 86 | + def set_context(**options) |
| 87 | + update_context(**options) |
| 88 | + yield if block_given? |
| 89 | + ensure |
| 90 | + update_context(**options.transform_values! { nil }) |
| 91 | + end |
| 92 | + |
| 93 | + # Temporarily tag any query executed within `&block`. Can be nested. |
| 94 | + def with_tag(tag, &block) |
| 95 | + inline_tags.push(tag) |
| 96 | + yield if block_given? |
| 97 | + ensure |
| 98 | + inline_tags.pop |
| 99 | + end |
| 100 | + |
| 101 | + def add_query_log_tags_to_sql(sql) # :nodoc: |
| 102 | + comments.each do |comment| |
| 103 | + unless sql.include?(comment) |
| 104 | + sql = prepend_comment ? "#{comment} #{sql}" : "#{sql} #{comment}" |
| 105 | + end |
| 106 | + end |
| 107 | + sql |
| 108 | + end |
| 109 | + |
| 110 | + private |
| 111 | + # Returns an array of comments which need to be added to the query, comprised |
| 112 | + # of configured and inline tags. |
| 113 | + def comments |
| 114 | + [ comment, inline_comment ].compact |
| 115 | + end |
| 116 | + |
| 117 | + # Returns an SQL comment +String+ containing the query log tags. |
| 118 | + # Sets and returns a cached comment if <tt>cache_query_log_tags</tt> is +true+. |
| 119 | + def comment |
| 120 | + if cache_query_log_tags |
| 121 | + self.cached_comment ||= uncached_comment |
| 122 | + else |
| 123 | + uncached_comment |
| 124 | + end |
| 125 | + end |
| 126 | + |
| 127 | + def uncached_comment |
| 128 | + content = tag_content |
| 129 | + if content.present? |
| 130 | + "/*#{escape_sql_comment(content)}*/" |
| 131 | + end |
| 132 | + end |
| 133 | + |
| 134 | + # Returns a +String+ containing any inline comments from +with_tag+. |
| 135 | + def inline_comment |
| 136 | + return nil unless inline_tags.present? |
| 137 | + "/*#{escape_sql_comment(inline_tag_content)}*/" |
| 138 | + end |
| 139 | + |
| 140 | + # Return the set of active inline tags from +with_tag+. |
| 141 | + def inline_tags |
| 142 | + context[:inline_tags] ||= [] |
| 143 | + end |
| 144 | + |
| 145 | + def context |
| 146 | + Thread.current[:active_record_query_log_tags_context] ||= {} |
| 147 | + end |
| 148 | + |
| 149 | + def escape_sql_comment(content) |
| 150 | + content.to_s.gsub(%r{ (/ (?: | \g<1>) \*) \+? \s* | \s* (\* (?: | \g<2>) /) }x, "") |
| 151 | + end |
| 152 | + |
| 153 | + def tag_content |
| 154 | + tags.flat_map { |i| [*i] }.filter_map do |tag| |
| 155 | + key, value_input = tag |
| 156 | + val = case value_input |
| 157 | + when nil then instance_exec(&taggings[key]) if taggings.has_key? key |
| 158 | + when Proc then instance_exec(&value_input) |
| 159 | + else value_input |
| 160 | + end |
| 161 | + "#{key}:#{val}" unless val.nil? |
| 162 | + end.join(",") |
| 163 | + end |
| 164 | + |
| 165 | + def inline_tag_content |
| 166 | + inline_tags.join |
| 167 | + end |
| 168 | + end |
| 169 | + |
| 170 | + module ExecutionMethods |
| 171 | + def execute(sql, *args, **kwargs) |
| 172 | + super(ActiveRecord::QueryLogs.add_query_log_tags_to_sql(sql), *args, **kwargs) |
| 173 | + end |
| 174 | + |
| 175 | + def exec_query(sql, *args, **kwargs) |
| 176 | + super(ActiveRecord::QueryLogs.add_query_log_tags_to_sql(sql), *args, **kwargs) |
| 177 | + end |
| 178 | + end |
| 179 | + end |
| 180 | +end |
| 181 | + |
| 182 | +ActiveSupport.on_load(:active_record) do |
| 183 | + ActiveRecord::QueryLogs.taggings.merge! \ |
| 184 | + socket: -> { ActiveRecord::Base.connection_db_config.socket }, |
| 185 | + db_host: -> { ActiveRecord::Base.connection_db_config.host }, |
| 186 | + database: -> { ActiveRecord::Base.connection_db_config.database } |
| 187 | +end |
0 commit comments