[DDC-2052] Add SqlWalker::setQueryComponent() to allow modification of the query component in a custom output walker

lib/Doctrine/ORM/Query/QueryException.php

@@ -153,4 +153,12 @@ class QueryException extends \Doctrine\ORM\ORMException
         return new self("Cannot check if a child of '" . $rootClass . "' is instanceof '" . $className . "', " .
                 "inheritance hierachy exists between these two classes.");
     }
+
+    public static function invalidQueryComponent($dqlAlias)
+    {
+        return new self(
+            "Invalid query component given for DQL alias '" . $dqlAlias . "', ".
+            "requires 'metadata', 'parent', 'relation', 'map', 'nestingLevel' and 'token' keys."
+        );
+    }
 }

lib/Doctrine/ORM/Query/SqlWalker.php

@@ -223,6 +223,23 @@ class SqlWalker implements TreeWalker
         return $this->queryComponents[$dqlAlias];
     }
 
+    /**
+     * Set or override a query component for a given dql alias.
+     *
+     * @param string $dqlAlias The DQL alias.
+     * @param array $queryComponent
+     */
+    protected function setQueryComponent($dqlAlias, array $queryComponent)
+    {
+        $requiredKeys = array('metadata', 'parent', 'relation', 'map', 'nestingLevel', 'token');
+
+        if (array_diff($requiredKeys, array_keys($queryComponent))) {
+            throw QueryException::invalidQueryComponent($dqlAlias);
+        }
+
+        $this->queryComponents[$dqlAlias] = $queryComponent;
+    }
+
     /**
      * Gets an executor that can be used to execute the result of this walker.
      *

tests/Doctrine/Tests/ORM/Query/CustomTreeWalkersTest.php

@@ -21,8 +21,6 @@ namespace Doctrine\Tests\ORM\Functional;
 
 use Doctrine\ORM\Query;
 
-require_once __DIR__ . '/../../TestInit.php';
-
 /**
  * Test case for custom AST walking and modification.
  *
@@ -40,18 +38,25 @@ class CustomTreeWalkersTest extends \Doctrine\Tests\OrmTestCase
         $this->_em = $this->_getTestEntityManager();
     }
 
-    public function assertSqlGeneration($dqlToBeTested, $sqlToBeConfirmed)
+    public function generateSql($dqlToBeTested, $treeWalkers, $outputWalker)
+    {
+        $query = $this->_em->createQuery($dqlToBeTested);
+        $query->setHint(Query::HINT_CUSTOM_TREE_WALKERS, $treeWalkers)
+              ->useQueryCache(false);
+
+        if ($outputWalker) {
+            $query->setHint(Query::HINT_CUSTOM_OUTPUT_WALKER, $outputWalker);
+        }
+
+        return $query->getSql();
+    }
+
+    public function assertSqlGeneration($dqlToBeTested, $sqlToBeConfirmed, $treeWalkers = array(), $outputWalker = null)
     {
         try {
-            $query = $this->_em->createQuery($dqlToBeTested);
-            $query->setHint(Query::HINT_CUSTOM_TREE_WALKERS, array('Doctrine\Tests\ORM\Functional\CustomTreeWalker'))
-                  ->useQueryCache(false);
-
-            $this->assertEquals($sqlToBeConfirmed, $query->getSql());
-            $query->free();
+            $this->assertEquals($sqlToBeConfirmed, $this->generateSql($dqlToBeTested, $treeWalkers, $outputWalker));
         } catch (\Exception $e) {
             $this->fail($e->getMessage() . ' at "' . $e->getFile() . '" on line ' . $e->getLine());
-
         }
     }
 
@@ -59,7 +64,8 @@ class CustomTreeWalkersTest extends \Doctrine\Tests\OrmTestCase
     {
         $this->assertSqlGeneration(
             'select u from Doctrine\Tests\Models\CMS\CmsUser u',
-            "SELECT c0_.id AS id0, c0_.status AS status1, c0_.username AS username2, c0_.name AS name3, c0_.email_id AS email_id4 FROM cms_users c0_ WHERE c0_.id = 1"
+            "SELECT c0_.id AS id0, c0_.status AS status1, c0_.username AS username2, c0_.name AS name3, c0_.email_id AS email_id4 FROM cms_users c0_ WHERE c0_.id = 1",
+            array('Doctrine\Tests\ORM\Functional\CustomTreeWalker')
         );
     }
 
@@ -67,7 +73,8 @@ class CustomTreeWalkersTest extends \Doctrine\Tests\OrmTestCase
     {
         $this->assertSqlGeneration(
             'select u from Doctrine\Tests\Models\CMS\CmsUser u where u.name = :name or u.name = :otherName',
-            "SELECT c0_.id AS id0, c0_.status AS status1, c0_.username AS username2, c0_.name AS name3, c0_.email_id AS email_id4 FROM cms_users c0_ WHERE (c0_.name = ? OR c0_.name = ?) AND c0_.id = 1"
+            "SELECT c0_.id AS id0, c0_.status AS status1, c0_.username AS username2, c0_.name AS name3, c0_.email_id AS email_id4 FROM cms_users c0_ WHERE (c0_.name = ? OR c0_.name = ?) AND c0_.id = 1",
+            array('Doctrine\Tests\ORM\Functional\CustomTreeWalker')
         );
     }
 
@@ -75,9 +82,30 @@ class CustomTreeWalkersTest extends \Doctrine\Tests\OrmTestCase
     {
         $this->assertSqlGeneration(
             'select u from Doctrine\Tests\Models\CMS\CmsUser u where u.name = :name',
-            "SELECT c0_.id AS id0, c0_.status AS status1, c0_.username AS username2, c0_.name AS name3, c0_.email_id AS email_id4 FROM cms_users c0_ WHERE c0_.name = ? AND c0_.id = 1"
+            "SELECT c0_.id AS id0, c0_.status AS status1, c0_.username AS username2, c0_.name AS name3, c0_.email_id AS email_id4 FROM cms_users c0_ WHERE c0_.name = ? AND c0_.id = 1",
+            array('Doctrine\Tests\ORM\Functional\CustomTreeWalker')
         );
     }
+
+    public function testSetUnknownQueryComponentThrowsException()
+    {
+        $this->setExpectedException("Doctrine\ORM\Query\QueryException", "Invalid query component given for DQL alias 'x', requires 'metadata', 'parent', 'relation', 'map', 'nestingLevel' and 'token' keys.");
+        $this->generateSql(
+            'select u from Doctrine\Tests\Models\CMS\CmsUser u',
+            array(),
+            __NAMESPACE__ . '\\AddUnknownQueryComponentWalker'
+        );
+    }
+}
+
+class AddUnknownQueryComponentWalker extends Query\SqlWalker
+{
+    public function walkSelectStatement(Query\AST\SelectStatement $selectStatement)
+    {
+        parent::walkSelectStatement($selectStatement);
+
+        $this->setQueryComponent('x', array());
+    }
 }
 
 class CustomTreeWalker extends Query\TreeWalkerAdapter