keeran commented on May 17 •

Porting Marginalia to Rails

This PR brings Marginalia SQL comments to Rails as a native feature. I have attempted to closely follow the Marginalia approach to solving the problem whilst introducing a new performance benefit for large applications.

The implementation (both in Marginalia and in this PR) can be broken down into 3 key areas:

Comment configuration & construction

Marginalia defines an array of components to construct an SQL comment. For each item in components, a method is called on Marginalia::Comment and the resulting key-value pair is concatenated into a comment string. This array of components can be modified by the application developer.

In Marginalia:

  self.components.each do |c|
    component_value = self.send(c)
    if component_value.present?
      ret << "#{c}:#{component_value},"
    end
  end

In this PR we have adjusted the method-calling approach to work with Procs instead. Default tags are defined as Procs in Railtie callback handlers (added to the taggings attribute), and users are able to pass custom tag definitions during application configuration.

  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(",")

If a custom comment component is required, a Hash can be added to the components array representing the key/value pair for the tag. If a Proc is passed as a value, it will be called when the comment is constructed and can reference internal state via the context Hash.

    config.active_record.query_log_tags_enabled = true
    config.active_record.query_log_tags = [
      :application,
      { single_custom: -> { "static_in_proc" } },
      {
        multiple: "new",
        values: "added",
        at: -> { Time.now }
      }
    ]

  # /*application:MyApp,single_custom:static_in_proc,multiple:new,values:added,at:2021-07-27 12:57:40 +0100*/

Rails integration - ActionController & ActiveJob

The core benefit of Marginalia is being able to decorate SQL queries with details of the context of the source of the query. By default Marginalia will add comment components representing the controller or job context and this is achieved by making use of callback filters.

Marginalia updates the controller context with around_* callbacks:

  if respond_to?(:around_action)
    around_action :record_query_comment
  else
    around_filter :record_query_comment
  end
...
  def record_query_comment
    Marginalia::Comment.update!(self)
    yield
  ensure
    Marginalia::Comment.clear!
  end

A similar approach is followed in this PR:

  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)

In order for these context-specific components to function, Marginalia maintains references to the controller & job in Thread.current:

  def self.marginalia_controller=(controller)
    Thread.current[:marginalia_controller] = controller
  end

  def self.marginalia_controller
    Thread.current[:marginalia_controller]
  end

Rather than create a potentially unbounded set of keys in Thread.current, I have opted to create a hash of context information (to include the controller and job) and only store that context hash in Thread.current, accessible via ActiveRecord::QueryLogs#update_context & ActiveRecord::QueryLogs#set_context.

Adapter SQL execution interception

Finally, when the SQL comment is ready to be inserted into the query, Marginalia performs a series of checks before alias_method chaining into the appropriate adapter methods.

Some of the complexity in this approach results from supporting older versions of Rails. Given we won’t have the same backwards-compatibility requirements here, I have simplified the instrumentation to prepend ing the execution methods on any subclass of AbstractAdapter:

  ActiveSupport.on_load(:active_record) do
    ConnectionAdapters::AbstractAdapter.descendants.each do |klass|
      klass.prepend(QueryLogs::ExecutionMethods) if klass.descendants.empty?
    end
  end
...
  module ExecutionMethods
    def execute(sql, *args, **kwargs)
      super(self.class.add_query_log_tags_to_sql(sql), *args, **kwargs)
    end

    def exec_query(sql, *args, **kwargs)
      super(self.class.add_query_log_tags_to_sql(sql), *args, **kwargs)
    end
  end

Notable changes

Comment caching

Introduces the ability to cache the component elements of a query comment to avoid needing to rebuild the same string repeatedly during a request / job if the context and content does not change.

ActiveRecord::QueryLogs.cache_query_log_tags = true

The cached comment is stored per thread and is reset whenever the context is updated via #update_context, #set_context, or via an explicit operation:

ActiveRecord::QueryLogs.cached_comment = nil

Migration

Rails 7 upgrade guide content in progress: migrating from Marginalia. (to be rewritten / adjusted as necessary when this feature has been completed)

/cc @eileencodes & @arthurnn