Add support for annotating queries generated by ActiveRecord::Relation with SQL comments

[mattyoho]

Summary

This pull request implements a small new feature that allows users to annotate generated SQL queries.

A new method, ActiveRecord::Relation#annotate, has been added that allows annotating queries generated from the relation instance with block-style SQL comments.

For example:

Post.annotate("this is a comment").where(id: 123).to_sql
# SELECT "posts".* FROM "posts" WHERE "posts"."id" = 123 /* this is a comment */

The hope is to land this feature in time for Rails 6.0. 🙂🤞

Discussion

This kind of annotation can be extremely useful when combined with instrumentation or other methods of tracking and analyzing queries generated at runtime, particularly in a large, mature application.

The marginalia gem provides similarly useful functionality in some contexts, but currently operates on all queries within a request or job and can't be used to target specific queries.

By hooking into ActiveRecord::Relation it's possible to annotate arbitrary scopes and model associations.

class Post < ActiveRecord::Base
  has_many :tags, -> { annotate("association comment") }

  scope :published, -> { where.not(published_at: nil).annotate("scope comment") }
end

This PR also updates ActiveRecord::Relation#update_all and ActiveRecord::Relation#delete_all to respect any annotations on the receiving relation object.

Multiple annotations may be added:

Post.where(id: 123).with_annotation("foo").annotate("bar")
# SELECT "posts".* FROM "posts" WHERE "posts"."id" = 123 /* foo */ /* bar */

Not all queries issued from ActiveRecord can be annotated using this approach, but many can.

Implementation

The implementation extends Arel under the hood by adding a new Arel::Nodes::Comment node and associated plumbing to work with them.

Of note, Arel::Nodes::Comment is responsible for preventing SQL injection via comment escaping.

Comment content is scrubbed of SQL block comment delimiters at creation time here. This should be reviewed for correctness. 👀

Example use with instrumentation

Here's a toy example of how this could be used with ActiveSupport::Notifications instrumentation to perform basic analysis.

query_checker = Proc.new do |_, _, _, _, payload|
  query = payload[:sql]
  if query =~ /JOIN "forbidden_table"/ && query !~ %r{/\* whitelisted \*/}
    Rails.logger.warn("this is bad!")
  end
end
ActiveSupport::Notifications.subscribe "sql.active_record", query_logger

# This would be all right
Post.annotate("whitelisted").joins(:forbidden_table).first
# But this would raise eyebrows
Post.joins(:forbidden_table).first

Rails Guides

It might be useful to add documentation for this new method to the Active Record Query Interface page.

I can provide a separate PR for that or add it here if so.

[jeremy]

Welcome addition and great, thorough PR. Thank you!

[mattyoho]

Thanks for the code review, @jeremy!

I've switched to adding distinct commits while I'm iterating based on review feedback. I can squash back down toward the end.

The Support Optimizer Hints PR in #35615 is really interesting in relation to this one. Edit: Or vice-versa. 😄

From #35615 (comment):

Should we ensure the hint is safe for a SQL comment, preventing query injection?

Could reuse the Arel Comment node #35617

I do think there'd be value in some slight generalization here to allow #35615 to reuse some of this. I've left some comments inline discussing it more.

In particular in this PR since the last review:

• 4a4003a renames Arel::Nodes::Comment to Arel::Nodes::Annotation.
• acdd316 demonstrates Arel::Nodes::Annotation representing an optimization hint within the visitor.

[mattyoho]

Howdy. 👋

I've iterated based on some review feedback and self-review, added some additional tests, and incorporated the latest changes from master, including #35615.

I think this patch is now in a state where it could be merged and so is ready for review again. 🙂🤞

I think it's worth me pointing out that I've refactored the additions in #35615 a bit after incorporating them into this branch. Hopefully that seems reasonable after looking at the diff.

In particular, I've changed Arel::Nodes::OptimizerHints to inherit from Arel::Nodes::Annotation rather than Arel::Nodes::Unary.

I think this makes conceptual sense, but the biggest practical advantage comes from reusing the comment sanitization routine rather than having two adjacent implementations.

You can see the changes specific to this refactor isolated in this Gist.

Thank you!

[mattyoho]

I've rebased off latest master and attempted to respond to the most recent feedback.

The #annotate query method now expects a String or Arel::Nodes::SqlLiteral argument.

Sanitization of SQL block comment delimiters is still happening at the node level, within the Arel::Nodes::Annotation mixin, for the reasons I proposed in #35617 (comment), but I'm of course willing to incorporate further feedback on that based on the discussion. 🙂

[mattyoho]

Sanitization of SQL block comment delimiters is still happening at the node level, within the Arel::Nodes::Annotation mixin, for the reasons I proposed in #35617 (comment), but I'm of course willing to incorporate further feedback on that based on the discussion.

When #35660 is merged, I'll update this diff to reuse that approach. 👍

[mattyoho]

The latest base branch changes have been incorporated and sanitization in the visitors has been embraced. I think this is ready for review again. /cc @jeremy

[eileencodes]

Looks good to me. @jeremy anymore concerns or do you think this is good to merge?

[mattyoho]

There is no worth to share Arel::Nodes::Annotation anymore.

Removed with latest push.

[kamipo]

How about joining multiple comments with "; "?

I think it is one solution to mitigate the issue Post.annotate("post comment").joins(:tags).merge(Tag.annotate("tag comment")) #=> "/* post comment tag comment */".

[mattyoho]

I don't think text content should be inserted into the user's annotation if it can be avoided. As a consumer of the API I wouldn't want that.

After more reflection I think the best approach by far is actually to wrap each annotation in its own delimiters, e.g.,

Post.annotate("post comment").joins(:tags).merge(Tag.annotate("tag comment"))
#=> "... /* post comment */ /* tag comment */"

This solves the concerns with running them together, it makes it much simpler to detect or parse out annotation in instrumentation, and it prevents different developers from stepping on one another's annotations.

I've been using a "backport" of this API in an experimental branch for a large refactor and after working with it a bit this is the behavior I'd want to have.

[mattyoho]

@kamipo Thank you for the thorough code review. The patch is much better because of it.

I believe this PR is once again ready for review and hopefully merging.

[arthurnn]

On marginalia we do some escaping on the comment, to make sure people will not annotate queries with user input and be open to SQL injections.
Is that necessary in here? (I scammed the code fast, and haven't saw any prevention like that)

[mattyoho]

@arthurnn That's being handled here:

rails/activerecord/lib/arel/visitors/to_sql.rb

Line 176 in f418258

collector << o.values.map { |v| "/* #{sanitize_as_sql_comment(v)} */" }.join(" ")

(Or here in this diff.)

[mattyoho]

(There are tests, e.g, here and here, as well.)

[abhaynikam]

Is annotate a finder method? Just curious to understand this because the documentation on the above line says

To retrieve objects from the database, Active Record provides several finder methods. Each finder method allows you to pass arguments into it to perform certain queries on your database without writing raw SQL.

As per the changelog, it says Add 'ActiveRecord::Relation#annotate' for adding SQL comments to its queries.

[mattyoho]

Hi, @abhaynikam! It's possible the description should be generalized a bit. Several other relation query methods arguably don't fit that documention sentence literally, either — preload, create_with, and includes being examples.

Perhaps the intro sentence could be updated in another PR.

[abhaynikam]

Thanks, @mattyoho . That makes sense, I agree a few other methods do not fit the description.

Great addition @mattyoho 🎉

[abhaynikam]

About updating the description, I would say let remove methods which do not fit description sinces that particular documentation section is for understanding how to retrieve objects from database.

I'll like to raise a PR if somebody from Rails team approve it is correct to update the doc 😊