Hardened fix update of inherited documents on findAndUpdate

lib/Doctrine/ODM/MongoDB/Persisters/DocumentPersister.php

@@ -871,7 +871,11 @@ public function addDiscriminatorToPreparedQuery(array $preparedQuery)
          */
         if ($this->class->hasDiscriminator() && ! isset($preparedQuery[$this->class->discriminatorField])) {
             $discriminatorValues = $this->getClassDiscriminatorValues($this->class);
-            $preparedQuery[$this->class->discriminatorField] = array('$in' => $discriminatorValues);
+            if ((count($discriminatorValues) === 1)) {
+                $preparedQuery[$this->class->discriminatorField] = $discriminatorValues[0];
+            } else {
+                $preparedQuery[$this->class->discriminatorField] = array('$in' => $discriminatorValues);
+            }
         }
 
         return $preparedQuery;

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

@@ -327,6 +327,12 @@ public function getQuery(array $options = array())
             $query['distinct'] = $documentPersister->prepareFieldName($query['distinct']);
         }
 
+        if ($this->class->inheritanceType === ClassMetadataInfo::INHERITANCE_TYPE_SINGLE_COLLECTION && ! empty($query['upsert']) &&
+            (empty($query['query'][$this->class->discriminatorField]) || is_array($query['query'][$this->class->discriminatorField]))) {
+            throw new \InvalidArgumentException('Upsert query that is to be performed on discriminated document does not have single ' .
+                'discriminator. Either not use base class or set \'' . $this->class->discriminatorField . '\' field manually.');
+        }
+
         if ( ! empty($query['select'])) {
             $query['select'] = $documentPersister->prepareSortOrProjection($query['select']);
             if ($this->hydrate && $this->class->inheritanceType === ClassMetadataInfo::INHERITANCE_TYPE_SINGLE_COLLECTION

tests/Doctrine/ODM/MongoDB/Tests/Functional/Ticket/GH971Test.php

@@ -0,0 +1,100 @@
+<?php
+
+
+namespace Doctrine\ODM\MongoDB\Tests\Functional\Ticket;
+
+use Doctrine\ODM\MongoDB\Mapping\Annotations as ODM;
+
+class GH971Test extends \Doctrine\ODM\MongoDB\Tests\BaseTest
+{
+    public function testUpdateOfInheritedDocumentUsingFindAndUpdate()
+    {
+        $name = "Ferrari";
+        $features = array(
+            "Super Engine",
+            "Huge Wheels"
+        );
+
+        //first query, create Car with name "Ferrari"
+        $this->dm->createQueryBuilder(__NAMESPACE__ . '\Car')
+            ->findAndUpdate()
+            ->upsert(true)
+            ->field('name')->equals($name)
+            ->sort('_id', -1)
+            ->field('features')->push($features[0])
+            ->getQuery()->execute();
+
+        //second query: update existing "Ferrari" with new feature
+        $this->dm->createQueryBuilder(__NAMESPACE__ . '\Car')
+            ->findAndUpdate()
+            ->upsert(true)
+            ->field('name')->equals($name)
+            ->sort('_id', -1)
+            ->field('features')->push($features[1])
+            ->getQuery()->execute();
+
+        $results = $this->dm->getRepository(__NAMESPACE__ . '\Car')->findAll();
+        $this->assertCount(1, $results);
+    }
+
+    /**
+     * @expectedException \InvalidArgumentException
+     * @expectedExceptionMessage Upsert query that is to be performed on discriminated document does not have single discriminator. Either not use base class or set 'type' field manually.
+     */
+    public function testUpsertThrowsExceptionWithIndecisiveDiscriminator()
+    {
+        $this->dm->createQueryBuilder(__NAMESPACE__ . '\Bicycle')
+            ->findAndUpdate()
+            ->upsert(true)
+            ->field('name')->equals("Cool")
+            ->field('features')->push("2 people")
+            ->getQuery()->execute();
+    }
+
+    public function testUpsertWillUseProvidedDiscriminator()
+    {
+        $this->dm->createQueryBuilder(__NAMESPACE__ . '\Bicycle')
+            ->findAndUpdate()
+            ->upsert(true)
+            ->field('type')->equals('tandem')
+            ->field('name')->equals("Cool")
+            ->field('features')->push("2 people")
+            ->getQuery()->execute();
+
+        $results = $this->dm->getRepository(__NAMESPACE__ . '\Tandem')->findAll();
+        $this->assertCount(1, $results);
+    }
+}
+
+/**
+ * @ODM\Document
+ * @ODM\InheritanceType("SINGLE_COLLECTION")
+ * @ODM\DiscriminatorField(fieldName="type")
+ * @ODM\DiscriminatorMap({"car"="Car", "bicycle"="Bicycle", "tandem"="Tandem"})
+ */
+class Vehicle
+{
+    /** @ODM\Id */
+    public $id;
+
+    /** @ODM\String */
+    public $name;
+
+    /** @ODM\EmbedMany */
+    public $features;
+}
+
+/**
+ * @ODM\Document
+ */
+class Car extends Vehicle {}
+
+/**
+ * @ODM\Document
+ */
+class Bicycle extends Vehicle {}
+
+/**
+ * @ODM\Document
+ */
+class Tandem extends Bicycle {}