From b0ac66549b5bcb1cb55e1e8d05a0f42cb29cf1b7 Mon Sep 17 00:00:00 2001
From: Vlasta Neubauer <cz.paranoiq@gmail.com>
Date: Mon, 8 Jul 2024 16:14:05 +0200
Subject: [PATCH] Fix subselect support (table aliases might not be resolved
 when handler is called and wrong alias is generated)

---
 src/Doctrine/MySql/UseIndexHintHandler.php    | 20 +++-----------
 .../Doctrine/MySql/UseIndexSqlWalkerTest.php  | 26 +++++++++++++++----
 2 files changed, 24 insertions(+), 22 deletions(-)

diff --git a/src/Doctrine/MySql/UseIndexHintHandler.php b/src/Doctrine/MySql/UseIndexHintHandler.php
index c65bc1e..5ee3b96 100644
--- a/src/Doctrine/MySql/UseIndexHintHandler.php
+++ b/src/Doctrine/MySql/UseIndexHintHandler.php
@@ -2,9 +2,7 @@
 
 namespace ShipMonk\Doctrine\MySql;
 
-use Doctrine\ORM\Query;
 use Doctrine\ORM\Query\AST\SelectStatement;
-use Doctrine\ORM\Query\Parser;
 use LogicException;
 use ShipMonk\Doctrine\Walker\HintHandler;
 use ShipMonk\Doctrine\Walker\SqlNode;
@@ -27,7 +25,7 @@ class UseIndexHintHandler extends HintHandler
      */
     public function getNodes(): array
     {
-        return [SqlNode::FromClause, SqlNode::SubselectFromClause];
+        return [SqlNode::SelectStatement, SqlNode::UpdateStatement, SqlNode::DeleteStatement];
     }
 
     public function processNode(SqlNode $sqlNode, string $sql): string
@@ -36,18 +34,6 @@ public function processNode(SqlNode $sqlNode, string $sql): string
         $sqlWalker = $this->getDoctrineSqlWalker();
         $query = $sqlWalker->getQuery();
         $platform = $query->getEntityManager()->getConnection()->getDatabasePlatform();
-        $dql = $query->getDQL();
-
-        if ($dql === null) {
-            throw new LogicException('Empty DQL query encountered');
-        }
-
-        $queryCopy = new Query($this->getDoctrineSqlWalker()->getEntityManager());
-        $queryCopy->setDQL($dql);
-        $parser = new Parser($queryCopy);
-
-        /** @var string $wholeSql */
-        $wholeSql = $parser->parse()->getSqlExecutor()->getSqlStatements();
 
         if (!is_a($platform, 'Doctrine\DBAL\Platforms\MySqlPlatform')) { // bypass platform MySqlPlatform => MySQLPlatform rename in dbal
             throw new LogicException("Only MySQL platform is supported, {$platform->getName()} given");
@@ -80,12 +66,12 @@ public function processNode(SqlNode $sqlNode, string $sql): string
                 : '\S+'; // doctrine always adds some alias
             $tableWithAliasRegex = "{$delimiter}{$tableName}\s+{$tableAlias}{$delimiter}i";
 
-            if (preg_match($tableWithAliasRegex, $wholeSql) === 0) {
+            if (preg_match($tableWithAliasRegex, $sql) === 0) {
                 $aliasInfo = $hint->getDqlAlias() !== null ? " with DQL alias {$hint->getDqlAlias()}" : '';
                 throw new LogicException("Invalid hint for index {$hint->getIndexName()}, table {$tableName}{$aliasInfo} is not present in the query.");
             }
 
-            if ($hint->getDqlAlias() === null && preg_match_all($tableWithAliasRegex, $wholeSql) !== 1) {
+            if ($hint->getDqlAlias() === null && preg_match_all($tableWithAliasRegex, $sql) !== 1) {
                 throw new LogicException("Invalid hint for index {$hint->getIndexName()}, table {$tableName} is present multiple times in the query, please specify DQL alias to apply index on a proper place.");
             }
 
diff --git a/tests/Doctrine/MySql/UseIndexSqlWalkerTest.php b/tests/Doctrine/MySql/UseIndexSqlWalkerTest.php
index bcb8f87..870563a 100644
--- a/tests/Doctrine/MySql/UseIndexSqlWalkerTest.php
+++ b/tests/Doctrine/MySql/UseIndexSqlWalkerTest.php
@@ -48,7 +48,10 @@ public function testWalker(string $dql, callable $configureQueryCallback, ?strin
     public static function walksProvider(): iterable
     {
         $userSelectDql = sprintf('SELECT u FROM %s u', User::class);
-        $userSubselectDql = sprintf('SELECT u FROM %s u WHERE u.id = (SELECT u2.id FROM %s u2 WHERE u2.id = 1)', User::class, User::class);
+        $userSelectWithJoinsDql = sprintf('SELECT u FROM %s u JOIN u.account a JOIN u.managedAccounts ma', User::class);
+        $userSubselectInWhereDql = sprintf('SELECT u FROM %s u WHERE u.id = (SELECT u2.id FROM %s u2 WHERE u2.id = 1)', User::class, User::class);
+        $userSubselectInSelectDql = sprintf('SELECT u, (SELECT u2.id FROM %s u2 WHERE u2.id = 1) as uu FROM %s u', User::class, User::class);
+        $userSubselectInSelectDql2 = sprintf('SELECT (SELECT u2.id FROM %s u2 WHERE u2.id = 1) as uu FROM %s u', User::class, User::class);
 
         yield 'FROM - use single index' => [
             $userSelectDql,
@@ -101,9 +104,6 @@ static function (Query $query): void {
             },
             'SELECT u0_.id AS id_0, u0_.account_id AS account_id_1 FROM user u0_ IGNORE INDEX (IDX_FOO, IDX_BAR)',
         ];
-
-        $userSelectWithJoinsDql = sprintf('SELECT u FROM %s u JOIN u.account a JOIN u.managedAccounts ma', User::class);
-
         yield 'JOIN - one single use index' => [
             $userSelectWithJoinsDql,
             static function (Query $query): void {
@@ -183,13 +183,29 @@ static function (Query $query): void {
         ];
 
         yield 'FROM in subselect' => [
-            $userSubselectDql,
+            $userSubselectInWhereDql,
             static function (Query $query): void {
                 $query->setHint(UseIndexHintHandler::class, [IndexHint::use('IDX_FOO', User::TABLE_NAME, 'u2')]);
             },
             'SELECT u0_.id AS id_0, u0_.account_id AS account_id_1 FROM user u0_ WHERE u0_.id = (SELECT u1_.id FROM user u1_ USE INDEX (IDX_FOO) WHERE u1_.id = 1)',
         ];
 
+        yield 'subselect and hint in main query' => [
+            $userSubselectInSelectDql,
+            static function (Query $query): void {
+                $query->setHint(UseIndexHintHandler::class, [IndexHint::use('IDX_FOO', User::TABLE_NAME, 'u')]);
+            },
+            'SELECT u0_.id AS id_0, (SELECT u1_.id FROM user u1_ WHERE u1_.id = 1) AS sclr_1, u0_.account_id AS account_id_2 FROM user u0_ USE INDEX (IDX_FOO)',
+        ];
+
+        yield 'subselect and hint in main query 2' => [
+            $userSubselectInSelectDql2,
+            static function (Query $query): void {
+                $query->setHint(UseIndexHintHandler::class, [IndexHint::use('IDX_FOO', User::TABLE_NAME, 'u')]);
+            },
+            'SELECT (SELECT u0_.id FROM user u0_ WHERE u0_.id = 1) AS sclr_0 FROM user u1_ USE INDEX (IDX_FOO)',
+        ];
+
         yield 'no hint' => [
             $userSelectDql,
             static function (Query $query): void {