Add QueryLogTags to Rails

Co-authored-by: Eileen M. Uchitelle <[email protected]>
Co-authored-by: Kasper Timm Hansen <[email protected]>

actionpack/lib/action_controller/railtie.rb

@@ -86,5 +86,35 @@ class Railtie < Rails::Railtie # :nodoc:
         ActionController::Metal.descendants.each(&:action_methods) if config.eager_load
       end
     end
+
+    initializer "action_controller.query_log_tags" do |app|
+      ActiveSupport.on_load(:action_controller_base) do
+        singleton_class.attr_accessor :log_query_tags_around_actions
+        self.log_query_tags_around_actions = true
+      end
+
+      ActiveSupport.on_load(:active_record) do
+        if app.config.active_record.query_log_tags_enabled && app.config.action_controller.log_query_tags_around_actions != false
+          ActiveRecord::QueryLogs.taggings.merge! \
+            controller:            -> { context[:controller]&.controller_name },
+            action:                -> { context[:controller]&.action_name },
+            namespaced_controller: -> { context[:controller]&.class&.name }
+
+          ActiveRecord::QueryLogs.tags << :controller << :action
+
+          context_extension = ->(controller) do
+            around_action :expose_controller_to_query_logs
+
+            private
+            def expose_controller_to_query_logs(&block)
+              ActiveRecord::QueryLogs.set_context(controller: self, &block)
+            end
+          end
+
+          ActionController::Base.class_eval(&context_extension)
+          ActionController::API.class_eval(&context_extension)
+        end
+      end
+    end
   end
 end

activejob/lib/active_job/railtie.rb

@@ -49,5 +49,28 @@ class Railtie < Rails::Railtie # :nodoc:
         end
       end
     end
+
+    initializer "active_job.query_log_tags" do |app|
+      ActiveSupport.on_load(:active_job) do
+        singleton_class.attr_accessor :log_query_tags_around_perform
+        self.log_query_tags_around_perform = true
+      end
+
+      ActiveSupport.on_load(:active_record) do
+        if app.config.active_record.query_log_tags_enabled && app.config.active_job.log_query_tags_around_perform != false
+          ActiveRecord::QueryLogs.taggings[:job] = -> { context[:job]&.class&.name }
+          ActiveRecord::QueryLogs.tags << :job
+
+          ActiveJob::Base.class_eval do
+            around_perform :expose_job_to_query_logs
+
+            private
+              def expose_job_to_query_logs(&block)
+                ActiveRecord::QueryLogs.set_context(job: self, &block)
+              end
+          end
+        end
+      end
+    end
   end
 end

activerecord/CHANGELOG.md

@@ -1,3 +1,47 @@
+*   Add `ActiveRecord::QueryLogs`.
+
+    Configurable tags can be automatically added to all SQL queries generated by Active Record.
+
+    ```ruby
+    # config/application.rb
+    module MyApp
+      class Application < Rails::Application
+        config.active_record.query_log_tags_enabled = true
+      end
+    end
+    ```
+
+    By default the application, controller and action details are added to the query tags:
+
+    ```ruby
+    class BooksController < ApplicationController
+      def index
+        @books = Book.all
+      end
+    end
+    ```
+
+    ```ruby
+    GET /books
+    # SELECT * FROM books /*application:MyApp;controller:books;action:index*/
+    ```
+
+    Custom tags containing static values and Procs can be defined in the application configuration:
+
+    ```ruby
+    config.active_record.query_log_tags = [
+      :application,
+      :controller,
+      :action,
+      {
+        custom_static: "foo",
+        custom_dynamic: -> { Time.now }
+      }
+    ]
+    ```
+
+    *Keeran Raj Hawoldar*, *Eileen M. Uchitelle*, *Kasper Timm Hansen*
+
 *   Added support for multiple databases to `rails db:setup` and `rails db:reset`.
 
     *Ryan Hall*

activerecord/lib/active_record.rb

@@ -58,6 +58,7 @@ module ActiveRecord
   autoload :Persistence
   autoload :QueryCache
   autoload :Querying
+  autoload :QueryLogs
   autoload :ReadonlyAttributes
   autoload :RecordInvalid, "active_record/validations"
   autoload :Reflection
@@ -322,6 +323,51 @@ def self.global_executor_concurrency # :nodoc:
   singleton_class.attr_accessor :verify_foreign_keys_for_fixtures
   self.verify_foreign_keys_for_fixtures = false
 
+  ##
+  # :singleton-method:
+  # Specify whether or not to enable adapter-level query comments.
+  # To enable:
+  #
+  #    config.active_record.query_log_tags_enabled = true
+  #
+  # When included in +ActionController+, controller context is automatically updated via an
+  # +around_action+ filter. This behaviour can be disabled as follows:
+  #
+  #    config.action_controller.log_query_tags_around_actions = false
+  #
+  # This behaviour can be disabled for +ActiveJob+ in a similar way:
+  #
+  #    config.active_job.log_query_tags_around_perform = false
+  singleton_class.attr_accessor :query_log_tags_enabled
+  self.query_log_tags_enabled = false
+
+  ##
+  # :singleton-method:
+  # An +Array+ specifying the key/value tags to be inserted in an SQL comment. Defaults to `[ :application ]`, a
+  # predefined tag returning the application name.
+  #
+  # Custom values can be passed in as a +Hash+:
+  #
+  #    config.active_record.query_log_tags = [ :application, { custom: 'value' } ]
+  #
+  # See +ActiveRecord::QueryLogs+ for more details
+  # on predefined tags and defining new tag content.
+  singleton_class.attr_accessor :query_log_tags
+  self.query_log_tags = [ :application ]
+
+  ##
+  # :singleton-method:
+  # Specify whether or not to enable caching of query log tags.
+  # For applications that have a large number of queries, caching query log tags can
+  # provide a performance benefit when the context does not change during the lifetime
+  # of the request or job execution.
+  #
+  # To enable:
+  #
+  #   config.active_record.cache_query_log_tags = true
+  singleton_class.attr_accessor :cache_query_log_tags
+  self.cache_query_log_tags = false
+
   def self.eager_load!
     super
     ActiveRecord::Locking.eager_load!

activerecord/lib/active_record/database_configurations/hash_config.rb

@@ -55,6 +55,10 @@ def host
         configuration_hash[:host]
       end
 
+      def socket # :nodoc:
+        configuration_hash[:socket]
+      end
+
       def database
         configuration_hash[:database]
       end

activerecord/lib/active_record/query_logs.rb

@@ -0,0 +1,187 @@
+# frozen_string_literal: true
+
+require "active_support/core_ext/module/attribute_accessors_per_thread"
+
+module ActiveRecord
+  # = Active Record Query Logs
+  #
+  # Automatically tag SQL queries with runtime information.
+  #
+  # Default tags available for use:
+  #
+  # * +application+
+  # * +pid+
+  # * +socket+
+  # * +db_host+
+  # * +database+
+  #
+  # _Action Controller and Active Job tags are also defined when used in Rails:_
+  #
+  # * +controller+
+  # * +action+
+  # * +job+
+  #
+  # The tags used in a query can be configured directly:
+  #
+  #     ActiveRecord::QueryLogs.tags = [ :application, :controller, :action, :job ]
+  #
+  # or via Rails configuration:
+  #
+  #     config.active_record.query_log_tags = [ :application, :controller, :action, :job ]
+  #
+  # To add new comment tags, add a hash to the tags array containing the keys and values you
+  # want to add to the comment. Dynamic content can be created by setting a proc or lambda value in a hash,
+  # and can reference any value stored in the +context+ object.
+  #
+  # Example:
+  #
+  #    tags = [
+  #      :application,
+  #      { custom_tag: -> { context[:controller].controller_name } }
+  #    ]
+  #    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)
+  #
+  # Temporary updates limited to the execution of a block:
+  #
+  #    ActiveRecord::QueryLogs.set_context(foo: Bar.new) do
+  #      posts = Post.all
+  #    end
+  #
+  # Tag comments can be prepended to the query:
+  #
+  #    ActiveRecord::QueryLogs.prepend_comment = true
+  #
+  # For applications where the content will not change during the lifetime of
+  # the request or job execution, the tags can be cached for reuse in every query:
+  #
+  #    ActiveRecord::QueryLogs.cache_query_log_tags = true
+  #
+  # This option can be set during application configuration or in a Rails initializer:
+  #
+  #    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
+    thread_mattr_accessor :cached_comment, instance_accessor: false
+
+    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 provided values to nil after
+      # the block is executed.
+      def set_context(**options)
+        update_context(**options)
+        yield if block_given?
+      ensure
+        update_context(**options.transform_values! { nil })
+      end
+
+      # Temporarily tag any query executed within `&block`. Can be nested.
+      def with_tag(tag, &block)
+        inline_tags.push(tag)
+        yield if block_given?
+      ensure
+        inline_tags.pop
+      end
+
+      def add_query_log_tags_to_sql(sql) # :nodoc:
+        comments.each do |comment|
+          unless sql.include?(comment)
+            sql = prepend_comment ? "#{comment} #{sql}" : "#{sql} #{comment}"
+          end
+        end
+        sql
+      end
+
+      private
+        # Returns an array of comments which need to be added to the query, comprised
+        # of configured and inline tags.
+        def comments
+          [ comment, inline_comment ].compact
+        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
+          if cache_query_log_tags
+            self.cached_comment ||= uncached_comment
+          else
+            uncached_comment
+          end
+        end
+
+        def uncached_comment
+          content = tag_content
+          if content.present?
+            "/*#{escape_sql_comment(content)}*/"
+          end
+        end
+
+        # Returns a +String+ containing any inline comments from +with_tag+.
+        def inline_comment
+          return nil unless inline_tags.present?
+          "/*#{escape_sql_comment(inline_tag_content)}*/"
+        end
+
+        # Return the set of active inline tags from +with_tag+.
+        def inline_tags
+          context[:inline_tags] ||= []
+        end
+
+        def context
+          Thread.current[:active_record_query_log_tags_context] ||= {}
+        end
+
+        def escape_sql_comment(content)
+          content.to_s.gsub(%r{ (/ (?: | \g<1>) \*) \+? \s* | \s* (\* (?: | \g<2>) /) }x, "")
+        end
+
+        def tag_content
+          tags.flat_map { |i| [*i] }.filter_map do |tag|
+            key, value_input = tag
+            val = case value_input
+                  when nil then instance_exec(&taggings[key]) if taggings.has_key? key
+                  when Proc then instance_exec(&value_input)
+                  else value_input
+            end
+            "#{key}:#{val}" unless val.nil?
+          end.join(",")
+        end
+
+        def inline_tag_content
+          inline_tags.join
+        end
+    end
+
+    module ExecutionMethods
+      def execute(sql, *args, **kwargs)
+        super(ActiveRecord::QueryLogs.add_query_log_tags_to_sql(sql), *args, **kwargs)
+      end
+
+      def exec_query(sql, *args, **kwargs)
+        super(ActiveRecord::QueryLogs.add_query_log_tags_to_sql(sql), *args, **kwargs)
+      end
+    end
+  end
+end
+
+ActiveSupport.on_load(:active_record) do
+  ActiveRecord::QueryLogs.taggings.merge! \
+    socket:   -> { ActiveRecord::Base.connection_db_config.socket },
+    db_host:  -> { ActiveRecord::Base.connection_db_config.host },
+    database: -> { ActiveRecord::Base.connection_db_config.database }
+end