From 1f6b49d2364576d36c515c785b9c7a0e4b8c3bc8 Mon Sep 17 00:00:00 2001
From: Guilherme Blanco <guilhermeblanco@gmail.com>
Date: Thu, 9 Jun 2011 15:42:40 -0300
Subject: [PATCH] Added getRootEntities to QueryBuilder.

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

diff --git a/lib/Doctrine/ORM/QueryBuilder.php b/lib/Doctrine/ORM/QueryBuilder.php
index 10eaebbe2..b67a550d8 100644
--- a/lib/Doctrine/ORM/QueryBuilder.php
+++ b/lib/Doctrine/ORM/QueryBuilder.php
@@ -240,7 +240,7 @@ class QueryBuilder
     }
 
     /**
-     * Gets the root alias of the query. This is the first entity alias involved
+     * Gets the root aliases of the query. This is the entity aliases involved
      * in the construction of the query.
      *
      * <code>
@@ -251,7 +251,7 @@ class QueryBuilder
      *     $qb->getRootAliases(); // array('u')
      * </code>
      *
-     * @return string $rootAlias
+     * @return array $rootAliases
      */
     public function getRootAliases()
     {
@@ -272,6 +272,39 @@ class QueryBuilder
         return $aliases;
     }
 
+    /**
+     * Gets the root entities of the query. This is the entity aliases involved
+     * in the construction of the query.
+     *
+     * <code>
+     *     $qb = $em->createQueryBuilder()
+     *         ->select('u')
+     *         ->from('User', 'u');
+     *
+     *     $qb->getRootEntities(); // array('User')
+     * </code>
+     *
+     * @return array $rootEntities
+     */
+    public function getRootEntities()
+    {
+        $entities = array();
+        
+        foreach ($this->_dqlParts['from'] as &$fromClause) {
+            if (is_string($fromClause)) {
+                $spacePos = strrpos($fromClause, ' ');
+                $from     = substr($fromClause, 0, $spacePos);
+                $alias    = substr($fromClause, $spacePos + 1);
+
+                $fromClause = new Query\Expr\From($from, $alias);
+            }
+            
+            $entities[] = $fromClause->getFrom();
+        }
+        
+        return $entities;
+    }
+
     /**
      * Sets a query parameter for the query being constructed.
      *
diff --git a/tests/Doctrine/Tests/ORM/QueryBuilderTest.php b/tests/Doctrine/Tests/ORM/QueryBuilderTest.php
index cf570f52c..2cde45223 100644
--- a/tests/Doctrine/Tests/ORM/QueryBuilderTest.php
+++ b/tests/Doctrine/Tests/ORM/QueryBuilderTest.php
@@ -652,6 +652,15 @@ class QueryBuilderTest extends \Doctrine\Tests\OrmTestCase
         $this->assertEquals(array('u'), $qb->getRootAliases());
     }
     
+    public function testGetRootEntities()
+    {
+        $qb = $this->_em->createQueryBuilder()
+            ->select('u')
+            ->from('Doctrine\Tests\Models\CMS\CmsUser', 'u');
+        
+        $this->assertEquals(array('Doctrine\Tests\Models\CMS\CmsUser'), $qb->getRootEntities());
+    }
+    
     public function testGetSeveralRootAliases()
     {
         $qb = $this->_em->createQueryBuilder()