feat: add withComment method to add comments

docs/en/reference/query-builder.rst

@@ -400,6 +400,49 @@ Multiple CTEs can be defined by calling the with method multiple times.
 
 Values of parameters used in a CTE should be defined in the main QueryBuilder.
 
+Comments
+~~~~~~~~
+
+To add comments to the query, use the ``addComment()`` method which
+will add the comment before the query:
+
+.. code-block:: php
+
+    <?php
+
+    $queryBuilder
+        ->select('id', 'name')
+        ->from('users')
+        ->addComment('This is a comment');
+    // /* This is a comment */ SELECT id, name FROM users
+
+Multiple comments can be added by calling the method multiple times:
+
+.. code-block:: php
+
+    <?php
+
+    $queryBuilder
+        ->select('id', 'name')
+        ->from('users')
+        ->addComment('Comment 1')
+        ->addComment('Comment 2');
+    // /* Comment 1 */ /* Comment 2 */ SELECT id, name FROM users
+
+Comments containing ``/*`` and ``*/`` will be sanitized in order to prevent
+comment injection. Each occurrence of aforementioned tokens will be replaced
+by an empty string and trimmed.
+
+.. code-block:: php
+
+    <?php
+
+    $queryBuilder
+        ->select('id', 'name')
+        ->from('users')
+        ->addComment('*/ drop table users; /*');
+    // /* drop table users; */ SELECT id, name FROM users
+
 Building Expressions
 --------------------
 

src/Query/QueryBuilder.php

@@ -27,8 +27,11 @@
 use function count;
 use function implode;
 use function is_object;
+use function preg_replace;
 use function sprintf;
+use function str_replace;
 use function substr;
+use function trim;
 
 /**
  * QueryBuilder class is responsible to dynamically create SQL queries.
@@ -170,6 +173,13 @@ class QueryBuilder
      */
     private array $commonTableExpressions = [];
 
+    /**
+     * Comments that will be added to the query.
+     *
+     * @var string[]
+     */
+    private array $comments = [];
+
     /**
      * The query cache profile used for caching results.
      */
@@ -358,7 +368,7 @@ public function executeStatement(): int|string
      */
     public function getSQL(): string
     {
-        return $this->sql ??= match ($this->type) {
+        return $this->sql ??= $this->getComments() . match ($this->type) {
             QueryType::INSERT => $this->getSQLForInsert(),
             QueryType::DELETE => $this->getSQLForDelete(),
             QueryType::UPDATE => $this->getSQLForUpdate(),
@@ -1625,4 +1635,37 @@ public function disableResultCache(): self
 
         return $this;
     }
+
+    public function addComment(string $comment): self
+    {
+        $sanitizedComment = $this->sanitizeComment($comment);
+        if ($sanitizedComment === '') {
+            return $this;
+        }
+
+        $this->comments[] = $sanitizedComment;
+
+        $this->sql = null;
+
+        return $this;
+    }
+
+    private function getComments(): string
+    {
+        $comments = '';
+        foreach ($this->comments as $comment) {
+            $comments .= sprintf('/* %s */ ', $comment);
+        }
+
+        return $comments;
+    }
+
+    private function sanitizeComment(string $comment): string
+    {
+        $comment = str_replace(['*/', '/*'], '', $comment);
+        $comment = preg_replace('/[[:cntrl:]]+/u', ' ', $comment) ?? $comment;
+        $comment = preg_replace('/\s+/u', ' ', $comment) ?? $comment;
+
+        return trim($comment);
+    }
 }

tests/Functional/Query/QueryBuilderTest.php

@@ -528,6 +528,26 @@ public function testSelectWithCTEUnion(): void
         self::assertSame($expectedRows, $qb->executeQuery()->fetchAllAssociative());
     }
 
+    public function testSelectAddCommentIsSanitized(): void
+    {
+        $expectedRows = $this->prepareExpectedRows([['id' => 1], ['id' => 2]]);
+        $qb           = $this->connection->createQueryBuilder();
+
+        $select = $qb
+            ->select('id')
+            ->from('for_update')
+            ->where('id IN (?, ?)')
+            ->setParameters([1, 2], [ParameterType::INTEGER, ParameterType::INTEGER])
+            ->addComment('Test comment')
+            ->addComment('*/ drop table users; /*');
+
+        self::assertSame(
+            '/* Test comment */ /* drop table users; */ SELECT id FROM for_update WHERE id IN (?, ?)',
+            $select->getSQL(),
+        );
+        self::assertSame($expectedRows, $select->executeQuery()->fetchAllAssociative());
+    }
+
     public function testPlatformDoesNotSupportCTE(): void
     {
         if ($this->platformSupportsCTEs()) {

tests/Query/QueryBuilderTest.php

@@ -65,6 +65,36 @@ public function testSimpleSelectWithoutFrom(): void
         self::assertEquals('SELECT some_function()', (string) $qb);
     }
 
+    public function testSimpleSelectAddComment(): void
+    {
+        $qb = new QueryBuilder($this->conn);
+
+        $qb->select('some_function()')
+           ->addComment('Test  comment');
+
+        self::assertEquals('/* Test comment */ SELECT some_function()', (string) $qb);
+    }
+
+    public function testEmptyCommentAfterSanitizationIsIgnored(): void
+    {
+        $qb = new QueryBuilder($this->conn);
+
+        $qb->addComment('/* */');
+        $qb->select('1');
+
+        self::assertSame('SELECT 1', (string) $qb);
+    }
+
+    public function testControlCharactersAreNormalizedToSpaces(): void
+    {
+        $qb = new QueryBuilder($this->conn);
+
+        $qb->addComment("route=read\nshard=2\ttrace=abc");
+        $qb->select('1');
+
+        self::assertSame('/* route=read shard=2 trace=abc */ SELECT 1', (string) $qb);
+    }
+
     public function testSimpleSelect(): void
     {
         $qb = new QueryBuilder($this->conn);
@@ -414,6 +444,17 @@ public function testUpdate(): void
         self::assertEquals('UPDATE users SET foo = ?, bar = ?', (string) $qb);
     }
 
+    public function testUpdateAddComment(): void
+    {
+        $qb = new QueryBuilder($this->conn);
+        $qb->update('users')
+           ->set('foo', '?')
+           ->set('bar', '?')
+           ->addComment('Test comment');
+
+        self::assertEquals('/* Test comment */ UPDATE users SET foo = ?, bar = ?', (string) $qb);
+    }
+
     public function testUpdateWhere(): void
     {
         $qb = new QueryBuilder($this->conn);
@@ -432,6 +473,15 @@ public function testDelete(): void
         self::assertEquals('DELETE FROM users', (string) $qb);
     }
 
+    public function testDeleteAddComment(): void
+    {
+        $qb = new QueryBuilder($this->conn);
+        $qb->delete('users');
+        $qb->addComment('Test comment');
+
+        self::assertEquals('/* Test comment */ DELETE FROM users', (string) $qb);
+    }
+
     public function testDeleteWhere(): void
     {
         $qb = new QueryBuilder($this->conn);
@@ -455,6 +505,21 @@ public function testInsertValues(): void
         self::assertEquals('INSERT INTO users (foo, bar) VALUES(?, ?)', (string) $qb);
     }
 
+    public function testInsertValuesAddComment(): void
+    {
+        $qb = new QueryBuilder($this->conn);
+        $qb->insert('users')
+            ->values(
+                [
+                    'foo' => '?',
+                    'bar' => '?',
+                ],
+            )
+            ->addComment('Test comment');
+
+        self::assertEquals('/* Test comment */ INSERT INTO users (foo, bar) VALUES(?, ?)', (string) $qb);
+    }
+
     public function testInsertReplaceValues(): void
     {
         $qb = new QueryBuilder($this->conn);