From 5ab4c3d50c47032a6295f3bb1aa240b3e159f3b6 Mon Sep 17 00:00:00 2001
From: Guilherme Blanco <guilhermeblanco@gmail.com>
Date: Wed, 25 Feb 2015 14:29:10 -0500
Subject: [PATCH] Added programatical support to define indexBy on root
 aliases.

---
 lib/Doctrine/ORM/QueryBuilder.php             | 44 +++++++++++++++++++
 tests/Doctrine/Tests/ORM/QueryBuilderTest.php | 30 +++++++++++++
 2 files changed, 74 insertions(+)

diff --git a/lib/Doctrine/ORM/QueryBuilder.php b/lib/Doctrine/ORM/QueryBuilder.php
index bb08514eb..23e0f328b 100644
--- a/lib/Doctrine/ORM/QueryBuilder.php
+++ b/lib/Doctrine/ORM/QueryBuilder.php
@@ -881,6 +881,50 @@ class QueryBuilder
         return $this->add('from', new Expr\From($from, $alias, $indexBy), true);
     }
 
+    /**
+     * Updates a query root corresponding to an entity setting its index by. This method is intended to be used with
+     * EntityRepository->createQueryBuilder(), which creates the initial FROM clause and do not allow you to update it
+     * setting an index by.
+     *
+     * <code>
+     *     $qb = $userRepository->createQueryBuilder('u')
+     *         ->indexBy('u', 'u.id');
+     *
+     *     // Is equivalent to...
+     *
+     *     $qb = $em->createQueryBuilder()
+     *         ->select('u')
+     *         ->from('User', 'u', 'u.id');
+     * </code>
+     *
+     * @param string $alias   The root alias of the class.
+     * @param string $indexBy The index for the from.
+     *
+     * @return QueryBuilder This QueryBuilder instance.
+     *
+     * @throws Query\QueryException
+     */
+    public function indexBy($alias, $indexBy)
+    {
+        $rootAliases = $this->getRootAliases();
+
+        if (!in_array($alias, $rootAliases)) {
+            throw new Query\QueryException(
+                sprintf('Specified root alias %s must be set before invoking indexBy().', $alias)
+            );
+        }
+
+        foreach ($this->_dqlParts['from'] as &$fromClause) {
+            if ($fromClause->getAlias() !== $alias) {
+                continue;
+            }
+
+            $fromClause = new Expr\From($fromClause->getFrom(), $fromClause->getAlias(), $indexBy);
+        }
+
+        return $this;
+    }
+
     /**
      * Creates and adds a join over an entity association to the query.
      *
diff --git a/tests/Doctrine/Tests/ORM/QueryBuilderTest.php b/tests/Doctrine/Tests/ORM/QueryBuilderTest.php
index d523f6f49..3c093bbbc 100644
--- a/tests/Doctrine/Tests/ORM/QueryBuilderTest.php
+++ b/tests/Doctrine/Tests/ORM/QueryBuilderTest.php
@@ -109,6 +109,25 @@ class QueryBuilderTest extends \Doctrine\Tests\OrmTestCase
         $this->assertValidQueryBuilder($qb, 'DELETE Doctrine\Tests\Models\CMS\CmsUser u');
     }
 
+    public function testSimpleSelectWithFromIndexBy()
+    {
+        $qb = $this->_em->createQueryBuilder()
+            ->from('Doctrine\Tests\Models\CMS\CmsUser', 'u', 'u.id')
+            ->select('u.id', 'u.username');
+
+        $this->assertValidQueryBuilder($qb, 'SELECT u.id, u.username FROM Doctrine\Tests\Models\CMS\CmsUser u INDEX BY u.id');
+    }
+
+    public function testSimpleSelectWithIndexBy()
+    {
+        $qb = $this->_em->createQueryBuilder()
+            ->from('Doctrine\Tests\Models\CMS\CmsUser', 'u')
+            ->indexBy('u', 'u.id')
+            ->select('u.id', 'u.username');
+
+        $this->assertValidQueryBuilder($qb, 'SELECT u.id, u.username FROM Doctrine\Tests\Models\CMS\CmsUser u INDEX BY u.id');
+    }
+
     public function testSimpleUpdate()
     {
         $qb = $this->_em->createQueryBuilder()
@@ -184,6 +203,17 @@ class QueryBuilderTest extends \Doctrine\Tests\OrmTestCase
         $this->assertValidQueryBuilder($qb, 'SELECT u, g FROM Doctrine\Tests\Models\CMS\CmsUser u, Doctrine\Tests\Models\CMS\CmsGroup g');
     }
 
+    public function testMultipleFromWithIndexBy()
+    {
+        $qb = $this->_em->createQueryBuilder()
+            ->select('u', 'g')
+            ->from('Doctrine\Tests\Models\CMS\CmsUser', 'u')
+            ->from('Doctrine\Tests\Models\CMS\CmsGroup', 'g')
+            ->indexBy('g', 'g.id');
+
+        $this->assertValidQueryBuilder($qb, 'SELECT u, g FROM Doctrine\Tests\Models\CMS\CmsUser u, Doctrine\Tests\Models\CMS\CmsGroup g INDEX BY g.id');
+    }
+
     public function testMultipleFromWithJoin()
     {
         $qb = $this->_em->createQueryBuilder()