Expr::references should go through discriminator map

lib/Doctrine/ODM/MongoDB/Mapping/MappingException.php

@@ -57,6 +57,27 @@ public static function mappingNotFound($className, $fieldName)
         return new self("No mapping found for field '$fieldName' in class '$className'.");
     }
 
+    /**
+     * @param string $className
+     * @param string $fieldName
+     * @return MappingException
+     */
+    public static function mappingNotFoundInClassNorDescendants($className, $fieldName)
+    {
+        return new self("No mapping found for field '$fieldName' in class '$className' nor its descendants.");
+    }
+
+    /**
+     * @param $fieldName
+     * @param $className
+     * @param $className2
+     * @return MappingException
+     */
+    public static function referenceFieldConflict($fieldName, $className, $className2)
+    {
+        return new self("Reference mapping for field '$fieldName' in class '$className' conflicts with one mapped in class '$className2'.");
+    }
+
     /**
      * @param string $document
      * @param string $fieldName

lib/Doctrine/ODM/MongoDB/Query/Expr.php

@@ -21,6 +21,7 @@
 
 use Doctrine\ODM\MongoDB\DocumentManager;
 use Doctrine\ODM\MongoDB\Mapping\ClassMetadata;
+use Doctrine\ODM\MongoDB\Mapping\MappingException;
 
 /**
  * Query expression builder for ODM.
@@ -60,7 +61,7 @@ public function setClassMetadata(ClassMetadata $class)
     public function references($document)
     {
         if ($this->currentField) {
-            $mapping = $this->class->getFieldMapping($this->currentField);
+            $mapping = $this->getReferenceMapping();
             $dbRef = $this->dm->createDBRef($document, $mapping);
 
             if (isset($mapping['simple']) && $mapping['simple']) {
@@ -90,7 +91,7 @@ public function references($document)
     public function includesReferenceTo($document)
     {
         if ($this->currentField) {
-            $mapping = $this->class->getFieldMapping($this->currentField);
+            $mapping = $this->getReferenceMapping();
             $dbRef = $this->dm->createDBRef($document, $mapping);
 
             if (isset($mapping['simple']) && $mapping['simple']) {
@@ -127,4 +128,37 @@ public function getNewObj()
             ->getDocumentPersister($this->class->name)
             ->prepareQueryOrNewObj($this->newObj);
     }
+
+    /**
+     * Gets reference mapping for current field from current class or its descendants.
+     *
+     * @return array
+     * @throws MappingException
+     */
+    private function getReferenceMapping()
+    {
+        $mapping = null;
+        try {
+            $mapping = $this->class->getFieldMapping($this->currentField);
+        } catch (MappingException $e) {
+            if (empty($this->class->discriminatorMap)) {
+                throw $e;
+            }
+            $foundIn = null;
+            foreach ($this->class->discriminatorMap as $child) {
+                $childClass = $this->dm->getClassMetadata($child);
+                if ($childClass->hasAssociation($this->currentField)) {
+                    if ($mapping !== null && $mapping !== $childClass->getFieldMapping($this->currentField)) {
+                        throw MappingException::referenceFieldConflict($this->currentField, $foundIn->name, $childClass->name);
+                    }
+                    $mapping = $childClass->getFieldMapping($this->currentField);
+                    $foundIn = $childClass;
+                }
+            }
+            if ($mapping === null) {
+                throw MappingException::mappingNotFoundInClassNorDescendants($this->class->name, $this->currentField);
+            }
+        }
+        return $mapping;
+    }
 }

tests/Doctrine/ODM/MongoDB/Tests/Query/BuilderTest.php

@@ -2,16 +2,157 @@
 
 namespace Doctrine\ODM\MongoDB\Tests\Query;
 
-use Doctrine\ODM\MongoDB\Query\Builder;
+use Doctrine\ODM\MongoDB\Mapping\Annotations as ODM;
+use Documents\Feature;
 
 class BuilderTest extends \Doctrine\ODM\MongoDB\Tests\BaseTest
 {
     /**
-     * @expectedException InvalidArgumentException
+     * @expectedException \InvalidArgumentException
      */
     public function testPrimeRequiresBooleanOrCallable()
     {
         $this->dm->createQueryBuilder('Documents\User')
             ->field('groups')->prime(1);
     }
+
+    public function testReferencesGoesThroughDiscriminatorMap()
+    {
+        $f = new Feature('Smarter references');
+        $this->dm->persist($f);
+
+        $q1 = $this->dm->createQueryBuilder(ParentClass::class)
+            ->field('featureFull')->references($f)
+            ->getQuery()->debug();
+
+        $this->assertEquals([ 'featureFull.$id' => new \MongoId($f->id) ], $q1['query']);
+
+        $q2 = $this->dm->createQueryBuilder(ParentClass::class)
+            ->field('featureSimple')->references($f)
+            ->getQuery()->debug();
+
+        $this->assertEquals([ 'featureSimple' => new \MongoId($f->id) ], $q2['query']);
+    }
+
+    /**
+     * @expectedException \Doctrine\ODM\MongoDB\Mapping\MappingException
+     * @expectedExceptionMessage No mapping found for field 'nope' in class 'Doctrine\ODM\MongoDB\Tests\Query\ParentClass' nor its descendants.
+     */
+    public function testReferencesThrowsSpecializedExceptionForDiscriminatedDocuments()
+    {
+        $f = new Feature('Smarter references');
+        $this->dm->persist($f);
+
+        $this->dm->createQueryBuilder(ParentClass::class)
+            ->field('nope')->references($f)
+            ->getQuery();
+    }
+
+    /**
+     * @expectedException \Doctrine\ODM\MongoDB\Mapping\MappingException
+     * @expectedExceptionMessage Reference mapping for field 'conflict' in class 'Doctrine\ODM\MongoDB\Tests\Query\ChildA' conflicts with one mapped in class 'Doctrine\ODM\MongoDB\Tests\Query\ChildB'.
+     */
+    public function testReferencesThrowsSpecializedExceptionForConflictingMappings()
+    {
+        $f = new Feature('Smarter references');
+        $this->dm->persist($f);
+
+        $this->dm->createQueryBuilder(ParentClass::class)
+            ->field('conflict')->references($f)
+            ->getQuery();
+    }
+
+    public function testIncludesReferenceToGoesThroughDiscriminatorMap()
+    {
+        $f = new Feature('Smarter references');
+        $this->dm->persist($f);
+
+        $q1 = $this->dm->createQueryBuilder(ParentClass::class)
+            ->field('featureFullMany')->includesReferenceTo($f)
+            ->getQuery()->debug();
+
+        $this->assertEquals([ 'featureFullMany' => [ '$elemMatch' => [ '$id' => new \MongoId($f->id) ] ] ], $q1['query']);
+
+        $q2 = $this->dm->createQueryBuilder(ParentClass::class)
+            ->field('featureSimpleMany')->includesReferenceTo($f)
+            ->getQuery()->debug();
+
+        $this->assertEquals([ 'featureSimpleMany' => [ '$elemMatch' => new \MongoId($f->id) ] ], $q2['query']);
+    }
+
+    /**
+     * @expectedException \Doctrine\ODM\MongoDB\Mapping\MappingException
+     * @expectedExceptionMessage No mapping found for field 'nope' in class 'Doctrine\ODM\MongoDB\Tests\Query\ParentClass' nor its descendants.
+     */
+    public function testIncludesReferenceToThrowsSpecializedExceptionForDiscriminatedDocuments()
+    {
+        $f = new Feature('Smarter references');
+        $this->dm->persist($f);
+
+        $this->dm->createQueryBuilder(ParentClass::class)
+            ->field('nope')->includesReferenceTo($f)
+            ->getQuery();
+    }
+
+    /**
+     * @expectedException \Doctrine\ODM\MongoDB\Mapping\MappingException
+     * @expectedExceptionMessage Reference mapping for field 'conflictMany' in class 'Doctrine\ODM\MongoDB\Tests\Query\ChildA' conflicts with one mapped in class 'Doctrine\ODM\MongoDB\Tests\Query\ChildB'.
+     */
+    public function testIncludesReferenceToThrowsSpecializedExceptionForConflictingMappings()
+    {
+        $f = new Feature('Smarter references');
+        $this->dm->persist($f);
+
+        $this->dm->createQueryBuilder(ParentClass::class)
+            ->field('conflictMany')->includesReferenceTo($f)
+            ->getQuery();
+    }
+}
+
+/**
+ * @ODM\Document
+ * @ODM\InheritanceType("SINGLE_COLLECTION")
+ * @ODM\DiscriminatorField(fieldName="type")
+ * @ODM\DiscriminatorMap({"ca"="ChildA", "cb"="ChildB"})
+ */
+class ParentClass
+{
+    /** @ODM\Id */
+    public $id;
+}
+
+/**
+ * @ODM\Document
+ */
+class ChildA extends ParentClass
+{
+    /** @ODM\ReferenceOne(targetDocument="Documents\Feature") */
+    public $featureFull;
+
+    /** @ODM\ReferenceMany(targetDocument="Documents\Feature") */
+    public $featureFullMany;
+
+    /** @ODM\ReferenceOne(targetDocument="Documents\Feature") */
+    public $conflict;
+
+    /** @ODM\ReferenceMany(targetDocument="Documents\Feature") */
+    public $conflictMany;
+}
+
+/**
+ * @ODM\Document
+ */
+class ChildB extends ParentClass
+{
+    /** @ODM\ReferenceOne(targetDocument="Documents\Feature", simple=true) */
+    public $featureSimple;
+
+    /** @ODM\ReferenceMany(targetDocument="Documents\Feature", simple=true) */
+    public $featureSimpleMany;
+
+    /** @ODM\ReferenceOne(targetDocument="Documents\Feature", simple=true) */
+    public $conflict;
+
+    /** @ODM\ReferenceMany(targetDocument="Documents\Feature", simple=true) */
+    public $conflictMany;
 }