Fixed #24638 -- Added QuerySet.comment() to inject SQL comments.

AUTHORS

@@ -980,6 +980,7 @@ answer newbie questions, and generally made Django that much better:
     Shannon -jj Behrens <https://www.jjinux.com/>
     Shawn Milochik <shawn@milochik.com>
     Shreya Bamne <shreya.bamne@gmail.com>
+    Sid <dcsid10@gmail.com>
     Silvan Spross <silvan.spross@gmail.com>
     Simeon Visser <http://simeonvisser.com>
     Simon Blanchard

django/db/models/query.py

@@ -39,6 +39,7 @@
 from django.utils import timezone
 from django.utils.deprecation import RemovedInDjango70Warning, django_file_prefixes
 from django.utils.functional import cached_property
+from django.utils.regex_helper import _lazy_re_compile
 
 # The maximum number of results to fetch in a get() query.
 MAX_GET_RESULTS = 21
@@ -47,6 +48,7 @@
 REPR_OUTPUT_SIZE = 20
 
 DEFAULT_FETCH_MODE = FETCH_ONE
+sql_comment_re = _lazy_re_compile(r"^[\w ._-]+$")
 
 
 class BaseIterable:
@@ -1705,6 +1707,9 @@ def _combinator_query(self, combinator, *other_qs, all=False):
         clone.query.default_ordering = True
         self._clear_ordering_in_combined_queries(clone.query, other_qs)
         clone.query.clear_limits()
+        # Comments belong to the combined leg (`self.query`). Clear them on
+        # the outer combinator wrapper so they aren't emitted twice.
+        clone.query.comments = ()
         clone.query.combinator = combinator
         clone.query.combinator_all = all
         return clone
@@ -1980,6 +1985,19 @@ def fetch_mode(self, fetch_mode):
         clone._fetch_mode = fetch_mode
         return clone
 
+    def comment(self, message):
+        """Add an SQL comment to the query."""
+        if not isinstance(message, str):
+            raise TypeError("QuerySet.comment() argument must be a string.")
+        if not sql_comment_re.fullmatch(message):
+            raise ValueError(
+                "QuerySet.comment() argument must contain only letters, numbers, "
+                "spaces, periods, underscores, and hyphens."
+            )
+        clone = self._chain()
+        clone.query.comments = (*clone.query.comments, message)
+        return clone
+
     ###################################
     # PUBLIC INTROSPECTION ATTRIBUTES #
     ###################################

django/db/models/sql/compiler.py

@@ -625,7 +625,10 @@ def get_combinator_sql(self, combinator, all):
         sql_parts, args_parts = zip(
             *((braces.format(sql), args) for sql, args in parts)
         )
-        result = [" {} ".format(combinator_sql).join(sql_parts)]
+        combined_sql = " {} ".format(combinator_sql).join(sql_parts)
+        if self.query.comments:
+            combined_sql = "%s %s" % (self.comments_sql(), combined_sql)
+        result = [combined_sql]
         params = []
         for part in args_parts:
             params.extend(part)
@@ -820,6 +823,9 @@ def as_sql(self, with_limits=True, with_col_aliases=False):
                     result += distinct_result
                     params += distinct_params
 
+                if self.query.comments:
+                    result.append(self.comments_sql())
+
                 out_cols = []
                 for _, (s_sql, s_params), alias in self.select + extra_select:
                     if alias:
@@ -1678,6 +1684,17 @@ def explain_query(self):
                 else:
                     yield value
 
+    def comments_sql(self):
+        """
+        Return SQL comments for ``Query.comments`` as a single space-separated
+        string of ``/* ... */`` blocks.
+
+        Validation in ``QuerySet.comment()`` restricts comments to word
+        characters, spaces, periods, underscores, and hyphens, so the values
+        here cannot contain SQL comment delimiters or control characters.
+        """
+        return " ".join("/* %s */" % comment for comment in self.query.comments)
+
 
 class SQLInsertCompiler(SQLCompiler):
     returning_fields = None
@@ -2089,10 +2106,10 @@ def as_sql(self):
                 values.append(f"{quoted_name} = %s")
                 update_params.append(val)
         table = self.query.base_table
-        result = [
-            "UPDATE %s SET" % qn(table),
-            ", ".join(values),
-        ]
+        result = ["UPDATE"]
+        if self.query.comments:
+            result.append(self.comments_sql())
+        result += [qn(table), "SET", ", ".join(values)]
         try:
             where, params = self.compile(self.query.where)
         except FullResultSet:

django/db/models/sql/query.py

@@ -302,6 +302,8 @@ class Query(BaseExpression):
 
     explain_info = None
 
+    comments = ()
+
     def __init__(self, model, alias_cols=True):
         self.model = model
         self.alias_refcount = {}

docs/ref/models/querysets.txt

@@ -1927,6 +1927,48 @@ For example:
     # queries the database with the 'backup' alias
     >>> Entry.objects.using("backup")
 
+``comment()``
+~~~~~~~~~~~~~
+
+.. versionadded:: 6.2
+
+.. method:: comment(message)
+
+Returns a new ``QuerySet`` that injects ``message`` into the resulting SQL as
+an ``/* ... */`` comment, after ``SELECT`` (and ``DISTINCT`` if present) or
+after ``UPDATE``.
+
+For example::
+
+    Entry.objects.comment("hello").all()
+
+roughly produces:
+
+.. code-block:: sql
+
+    SELECT /* hello */ "blog_entry"."id" FROM "blog_entry"
+
+Calls can be chained; each call appends a separate block::
+
+    Entry.objects.comment("first").comment("second").update(headline="x")
+
+produces:
+
+.. code-block:: sql
+
+    UPDATE /* first */ /* second */ "blog_entry" SET "headline" = 'x'
+
+``message`` must be a non-empty string containing only letters, numbers,
+spaces, periods, underscores, and hyphens. ``comment()`` raises
+:exc:`ValueError` if ``message`` contains any other characters, and
+:exc:`TypeError` if ``message`` is not a string.
+
+``comment()`` only affects ``SELECT`` and ``UPDATE`` statements; ``DELETE``,
+``INSERT``, and statements issued by migrations are not annotated. When
+called on a combined ``QuerySet`` (produced by :meth:`union`,
+:meth:`intersection`, or :meth:`difference`), the comment is emitted as a
+leading ``/* ... */`` block before the combined SQL.
+
 ``select_for_update()``
 ~~~~~~~~~~~~~~~~~~~~~~~
 

docs/releases/6.2.txt

@@ -183,7 +183,8 @@ Migrations
 Models
 ~~~~~~
 
-* ...
+* The new :meth:`.QuerySet.comment` method allows injecting an SQL comment
+  into the generated ``SELECT`` or ``UPDATE`` statement.
 
 Requests and Responses
 ~~~~~~~~~~~~~~~~~~~~~~

tests/basic/tests.py

@@ -815,6 +815,7 @@ class ManagerTest(SimpleTestCase):
         "aupdate",
         "aupdate_or_create",
         "fetch_mode",
+        "comment",
     ]
 
     def test_manager_methods(self):

tests/queries/tests.py

@@ -4641,6 +4641,138 @@ def test_ticket_23622(self):
         self.assertSequenceEqual(Ticket23605A.objects.filter(qx), [a2])
 
 
+class QuerySetCommentTests(TestCase):
+    """Tests for QuerySet.comment() (Trac #24638)."""
+
+    def test_select_single_comment(self):
+        with CaptureQueriesContext(connection) as ctx:
+            list(NamedCategory.objects.comment("blog.views.py 50"))
+        self.assertIn("SELECT /* blog.views.py 50 */ ", ctx.captured_queries[0]["sql"])
+
+    def test_select_multiple_comments_preserve_order(self):
+        with CaptureQueriesContext(connection) as ctx:
+            list(
+                NamedCategory.objects.comment("first")
+                .comment("second")
+                .comment("third")
+            )
+        self.assertIn(
+            "SELECT /* first */ /* second */ /* third */ ",
+            ctx.captured_queries[0]["sql"],
+        )
+
+    def test_select_with_distinct(self):
+        with CaptureQueriesContext(connection) as ctx:
+            list(NamedCategory.objects.distinct().comment("after distinct"))
+        sql = ctx.captured_queries[0]["sql"]
+        self.assertIn("DISTINCT", sql)
+        self.assertIn("/* after distinct */", sql)
+        # The comment must come after DISTINCT.
+        self.assertLess(sql.index("DISTINCT"), sql.index("/* after distinct */"))
+
+    def test_select_with_filter_preserves_comment(self):
+        with CaptureQueriesContext(connection) as ctx:
+            list(NamedCategory.objects.comment("request-id 123").filter(name="x"))
+        self.assertIn("/* request-id 123 */", ctx.captured_queries[0]["sql"])
+
+    def test_select_subquery(self):
+        with CaptureQueriesContext(connection) as ctx:
+            list(
+                NamedCategory.objects.annotate(
+                    foo=(
+                        Tag.objects.filter(category=OuterRef("id"))
+                        .comment("inner subquery")
+                        .values("name")[:1]
+                    )
+                ).comment("outer query")
+            )
+        sql = ctx.captured_queries[0]["sql"]
+        self.assertIn("/* inner subquery */", sql)
+        self.assertIn("/* outer query */", sql)
+
+    def test_aggregate_subquery(self):
+        with CaptureQueriesContext(connection) as ctx:
+            Tag.objects.comment("tag-aggregate").values("parent").annotate(
+                tag_per_parent=Count("pk")
+            ).aggregate(Max("tag_per_parent"))
+        sql = ctx.captured_queries[0]["sql"]
+        self.assertIn("/* tag-aggregate */", sql)
+
+    def test_update(self):
+        NamedCategory.objects.create(name="initial")
+        with CaptureQueriesContext(connection) as ctx:
+            NamedCategory.objects.comment("update from worker-3").update(name="renamed")
+        update_sql = next(
+            q["sql"] for q in ctx.captured_queries if q["sql"].startswith("UPDATE")
+        )
+        self.assertIn("UPDATE /* update from worker-3 */ ", update_sql)
+
+    def test_clone_preserves_comment(self):
+        base = NamedCategory.objects.comment("shared")
+        derived = base.filter(name="x").comment("extra")
+        with CaptureQueriesContext(connection) as ctx:
+            list(base)
+            list(derived)
+        self.assertIn("/* shared */", ctx.captured_queries[0]["sql"])
+        self.assertNotIn("/* extra */", ctx.captured_queries[0]["sql"])
+        derived_sql = ctx.captured_queries[1]["sql"]
+        self.assertIn("/* shared */", derived_sql)
+        self.assertIn("/* extra */", derived_sql)
+
+    def test_unicode_comment(self):
+        with CaptureQueriesContext(connection) as ctx:
+            list(NamedCategory.objects.comment("источник cron"))
+        self.assertIn("/* источник cron */", ctx.captured_queries[0]["sql"])
+
+    def test_comment_inside_union_leg(self):
+        qs = NamedCategory.objects.comment("first leg").union(
+            NamedCategory.objects.all()
+        )
+        with CaptureQueriesContext(connection) as ctx:
+            list(qs)
+        sql = ctx.captured_queries[0]["sql"]
+        self.assertEqual(sql.count("/* first leg */"), 1)
+
+    def test_comment_around_union(self):
+        qs = NamedCategory.objects.union(NamedCategory.objects.all()).comment(
+            "outer combinator"
+        )
+        with CaptureQueriesContext(connection) as ctx:
+            list(qs)
+        sql = ctx.captured_queries[0]["sql"]
+        self.assertIn("/* outer combinator */", sql)
+        self.assertLess(sql.index("/* outer combinator */"), sql.index("UNION"))
+
+    def test_rejects_non_string(self):
+        msg = "QuerySet.comment() argument must be a string."
+        for bad in [None, 42, b"bytes", ["list"], object()]:
+            with self.subTest(bad=bad), self.assertRaisesMessage(TypeError, msg):
+                NamedCategory.objects.comment(bad)
+
+    def test_rejects_disallowed_comment_characters(self):
+        msg = (
+            "QuerySet.comment() argument must contain only letters, numbers, "
+            "spaces, periods, underscores, and hyphens."
+        )
+        for bad in [
+            "",
+            "blog/views.py:50",
+            "request-id=123",
+            "foo /* bar",
+            "foo */ bar",
+            "*/-- DROP TABLE x;--/*",
+            "/* nested */",
+            "quote '",
+            'quote "',
+            "semi;colon",
+            "hash#tag",
+            "comma,value",
+            *(f"name{chr(c)}" for c in chain(range(32), range(0x7F, 0xA0))),
+        ]:
+            with self.subTest(bad=bad), self.assertRaisesMessage(ValueError, msg):
+                NamedCategory.objects.comment(bad)
+
+
 class QuerySetCloningTests(TestCase):
     @classmethod
     def setUpTestData(cls):