Fix field name translation in $lookup stage

doctrine · Oct 9, 2017 · 30b5b27 · 30b5b27
1 parent f5bc642
commit 30b5b27
Show file tree

Hide file tree

Showing 4 changed files with 113 additions and 5 deletions.
diff --git a/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/Lookup.php b/lib/Doctrine/ODM/MongoDB/Aggregation/Stage/Lookup.php
@@ -42,6 +42,11 @@ class Lookup extends BaseStage\Lookup
      */
     private $class;
 
+    /**
+     * @var ClassMetadata
+     */
+    private $targetClass;
+
     /**
      * @param Builder $builder
      * @param string $from
@@ -73,11 +78,16 @@ public function from($from)
 
         // Check if mapped class with given name exists
         try {
-            $targetMapping = $this->dm->getClassMetadata($from);
-            return parent::from($targetMapping->getCollection());
+            $this->targetClass = $this->dm->getClassMetadata($from);
         } catch (BaseMappingException $e) {
             return parent::from($from);
         }
+
+        if ($this->targetClass->isSharded()) {
+            throw MappingException::cannotUseShardedCollectionInLookupStages($this->targetClass->name);
+        }
+
+        return parent::from($this->targetClass->getCollection());
     }
 
     /**
@@ -92,8 +102,12 @@ private function fromReference($fieldName)
         }
 
         $referenceMapping = $this->class->getFieldMapping($fieldName);
-        $targetMapping = $this->dm->getClassMetadata($referenceMapping['targetDocument']);
-        parent::from($targetMapping->getCollection());
+        $this->targetClass = $this->dm->getClassMetadata($referenceMapping['targetDocument']);
+        if ($this->targetClass->isSharded()) {
+            throw MappingException::cannotUseShardedCollectionInLookupStages($this->targetClass->name);
+        }
+
+        parent::from($this->targetClass->getCollection());
 
         if ($referenceMapping['isOwningSide']) {
             switch ($referenceMapping['storeAs']) {
@@ -114,7 +128,7 @@ private function fromReference($fieldName)
                 throw MappingException::repositoryMethodLookupNotAllowed($this->class->name, $fieldName);
             }
 
-            $mappedByMapping = $targetMapping->getFieldMapping($referenceMapping['mappedBy']);
+            $mappedByMapping = $this->targetClass->getFieldMapping($referenceMapping['mappedBy']);
             switch ($mappedByMapping['storeAs']) {
                 case ClassMetadataInfo::REFERENCE_STORE_AS_ID:
                 case ClassMetadataInfo::REFERENCE_STORE_AS_REF:
@@ -132,4 +146,37 @@ private function fromReference($fieldName)
 
         return $this;
     }
+
+    /**
+     * {@inheritdoc}
+     */
+    public function localField($localField)
+    {
+        return parent::localField($this->prepareFieldName($localField, $this->class));
+    }
+
+    /**
+     * {@inheritdoc}
+     */
+    public function foreignField($foreignField)
+    {
+        return parent::foreignField($this->prepareFieldName($foreignField, $this->targetClass));
+    }
+
+    protected function prepareFieldName($fieldName, ClassMetadata $class = null)
+    {
+        if (!$class) {
+            return $fieldName;
+        }
+
+        return $this->getDocumentPersister($class)->prepareFieldName($fieldName);
+    }
+
+    /**
+     * @return \Doctrine\ODM\MongoDB\Persisters\DocumentPersister
+     */
+    private function getDocumentPersister(ClassMetadata $class)
+    {
+        return $this->dm->getUnitOfWork()->getDocumentPersister($class->name);
+    }
 }
diff --git a/lib/Doctrine/ODM/MongoDB/Mapping/MappingException.php b/lib/Doctrine/ODM/MongoDB/Mapping/MappingException.php
@@ -386,6 +386,15 @@ public static function cannotUseShardedCollectionInOutStage($className)
         return new self("Cannot use class '$className' as collection for out stage. Sharded collections are not allowed.");
     }
 
+    /**
+     * @param string $className
+     * @return MappingException
+     */
+    public static function cannotUseShardedCollectionInLookupStages($className)
+    {
+        return new self("Cannot use class '$className' as collection for lookup or graphLookup stage. Sharded collections are not allowed.");
+    }
+
     /**
      * @param string $className
      * @param string $fieldName

diff --git a/tests/Doctrine/ODM/MongoDB/Tests/Aggregation/Stage/LookupTest.php b/tests/Doctrine/ODM/MongoDB/Tests/Aggregation/Stage/LookupTest.php
@@ -2,6 +2,10 @@
 
 namespace Doctrine\ODM\MongoDB\Tests\Aggregation\Stage;
 
+use Doctrine\ODM\MongoDB\Mapping\MappingException;
+use Documents\CmsComment;
+use Documents\Sharded\ShardedOne;
+
 class LookupTest extends \Doctrine\ODM\MongoDB\Tests\BaseTest
 {
     /**
@@ -40,6 +44,29 @@ public function testLookupStage()
         $this->assertSame('alcaeus', $result[0]['user'][0]['username']);
     }
 
+    public function testLookupStageWithFieldNameTranslation()
+    {
+        $builder = $this->dm->createAggregationBuilder(\Documents\SimpleReferenceUser::class);
+        $builder
+            ->lookup(CmsComment::class)
+                ->localField('id')
+                ->foreignField('authorIp')
+                ->alias('user');
+
+        $expectedPipeline = [
+            [
+                '$lookup' => [
+                    'from' => 'CmsComment',
+                    'localField' => '_id',
+                    'foreignField' => 'ip',
+                    'as' => 'user',
+                ]
+            ]
+        ];
+
+        $this->assertEquals($expectedPipeline, $builder->getPipeline());
+    }
+
     public function testLookupStageWithClassName()
     {
         $builder = $this->dm->createAggregationBuilder(\Documents\SimpleReferenceUser::class);
@@ -349,6 +376,26 @@ public function testLookupStageReferenceManyInverseStoreAsRef()
         $this->assertCount(1, $result[0]['embeddedReferenceManyInverse']);
     }
 
+    public function testLookupToShardedCollectionThrowsException()
+    {
+        $builder = $this->dm->createAggregationBuilder(\Documents\User::class);
+
+        $this->expectException(MappingException::class);
+        $builder
+            ->lookup(ShardedOne::class)
+                ->localField('id')
+                ->foreignField('id');
+    }
+
+    public function testLookupToShardedReferenceThrowsException()
+    {
+        $builder = $this->dm->createAggregationBuilder(ShardedOne::class);
+
+        $this->expectException(MappingException::class);
+        $builder
+            ->lookup('user');
+    }
+
     private function insertTestData()
     {
         $user1 = new \Documents\User();

diff --git a/tests/Documents/Sharded/ShardedOne.php b/tests/Documents/Sharded/ShardedOne.php
@@ -18,4 +18,9 @@ class ShardedOne
 
     /** @ODM\Field(name="k", type="string") */
     public $key = 'testing';
+
+    /**
+     * @ODM\ReferenceOne(targetDocument=ShardedUser::class)
+     */
+    public $user;
 }