Merge pull request #1551 from malarzm/gh1538

Fix, test and document modifying references with Builder

docs/en/reference/query-builder-api.rst

@@ -664,7 +664,14 @@ method:
 Update Queries
 ~~~~~~~~~~~~~~
 
-Doctrine also supports executing atomic update queries using the `Query\Builder` object. You can use the conditional operations in combination with the ability to change document field values atomically. You have several modifier operations available to you that make it easy to update documents in Mongo:
+Doctrine also supports executing atomic update queries using the `Query\Builder`
+object. You can use the conditional operations in combination with the ability to
+change document field values atomically. Additionally if you are modifying a field
+that is a reference you can pass managed document to the Builder and let ODM build
+``DBRef`` object for you.
+
+You have several modifier operations
+available to you that make it easy to update documents in Mongo:
 
 * ``set($name, $value, $atomic = true)``
 * ``setNewObj($newObj)``

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

@@ -971,27 +971,28 @@ public function addFilterToPreparedQuery(array $preparedQuery)
      * PHP field names and types will be converted to those used by MongoDB.
      *
      * @param array $query
+     * @param bool $isNewObj
      * @return array
      */
-    public function prepareQueryOrNewObj(array $query)
+    public function prepareQueryOrNewObj(array $query, $isNewObj = false)
     {
         $preparedQuery = array();
 
         foreach ($query as $key => $value) {
             // Recursively prepare logical query clauses
             if (in_array($key, array('$and', '$or', '$nor')) && is_array($value)) {
                 foreach ($value as $k2 => $v2) {
-                    $preparedQuery[$key][$k2] = $this->prepareQueryOrNewObj($v2);
+                    $preparedQuery[$key][$k2] = $this->prepareQueryOrNewObj($v2, $isNewObj);
                 }
                 continue;
             }
 
             if (isset($key[0]) && $key[0] === '$' && is_array($value)) {
-                $preparedQuery[$key] = $this->prepareQueryOrNewObj($value);
+                $preparedQuery[$key] = $this->prepareQueryOrNewObj($value, $isNewObj);
                 continue;
             }
 
-            $preparedQueryElements = $this->prepareQueryElement($key, $value, null, true);
+            $preparedQueryElements = $this->prepareQueryElement($key, $value, null, true, $isNewObj);
             foreach ($preparedQueryElements as list($preparedKey, $preparedValue)) {
                 $preparedQuery[$preparedKey] = is_array($preparedValue)
                     ? array_map('\Doctrine\ODM\MongoDB\Types\Type::convertPHPToDatabaseValue', $preparedValue)
@@ -1011,10 +1012,11 @@ public function prepareQueryOrNewObj(array $query)
      * @param string $fieldName
      * @param mixed $value
      * @param ClassMetadata $class        Defaults to $this->class
-     * @param boolean $prepareValue Whether or not to prepare the value
+     * @param bool $prepareValue Whether or not to prepare the value
+     * @param bool $inNewObj Whether or not newObj is being prepared
      * @return array An array of tuples containing prepared field names and values
      */
-    private function prepareQueryElement($fieldName, $value = null, $class = null, $prepareValue = true)
+    private function prepareQueryElement($fieldName, $value = null, $class = null, $prepareValue = true, $inNewObj = false)
     {
         $class = isset($class) ? $class : $this->class;
 
@@ -1037,7 +1039,7 @@ private function prepareQueryElement($fieldName, $value = null, $class = null, $
 
             if (! empty($mapping['reference']) && is_object($value) && ! ($value instanceof \MongoId)) {
                 try {
-                    return $this->prepareDbRefElement($fieldName, $value, $mapping);
+                    return $this->prepareDbRefElement($fieldName, $value, $mapping, $inNewObj);
                 } catch (MappingException $e) {
                     // do nothing in case passed object is not mapped document
                 }
@@ -1400,11 +1402,15 @@ private function getWriteOptions(array $options = array())
      * @param string $fieldName
      * @param mixed $value
      * @param array $mapping
+     * @param bool $inNewObj
      * @return array
      */
-    private function prepareDbRefElement($fieldName, $value, array $mapping)
+    private function prepareDbRefElement($fieldName, $value, array $mapping, $inNewObj)
     {
         $dbRef = $this->dm->createDBRef($value, $mapping);
+        if ($inNewObj) {
+            return [[$fieldName, $dbRef]];
+        }
         $keys = ['$ref' => true, '$id' => true, '$db' => true];
         if ($mapping['storeAs'] === ClassMetadataInfo::REFERENCE_STORE_AS_ID) {
             unset($keys['$db']);

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

@@ -160,7 +160,7 @@ public function getNewObj()
     {
         return $this->dm->getUnitOfWork()
             ->getDocumentPersister($this->class->name)
-            ->prepareQueryOrNewObj($this->newObj);
+            ->prepareQueryOrNewObj($this->newObj, true);
     }
 
     /**

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

@@ -135,6 +135,54 @@ public function testIncludesReferenceToThrowsSpecializedExceptionForConflictingM
             ->field('conflictMany')->includesReferenceTo($f)
             ->getQuery();
     }
+
+    /**
+     * @dataProvider provideArrayUpdateOperatorsOnReferenceMany
+     */
+    public function testArrayUpdateOperatorsOnReferenceMany($class, $field)
+    {
+        $f = new Feature('Smarter references');
+        $this->dm->persist($f);
+
+        $q1 = $this->dm->createQueryBuilder($class)
+            ->findAndUpdate()
+            ->field($field)->addToSet($f)
+            ->getQuery()->debug();
+
+        $expected = $this->dm->createDBRef($f, $this->dm->getClassMetadata($class)->fieldMappings[$field]);
+        $this->assertEquals($expected, $q1['newObj']['$addToSet'][$field]);
+    }
+
+    public function provideArrayUpdateOperatorsOnReferenceMany()
+    {
+        yield [ChildA::class, 'featureFullMany'];
+        yield [ChildB::class, 'featureSimpleMany'];
+        yield [ChildC::class, 'featurePartialMany'];
+    }
+
+    /**
+     * @dataProvider provideArrayUpdateOperatorsOnReferenceOne
+     */
+    public function testArrayUpdateOperatorsOnReferenceOne($class, $field)
+    {
+        $f = new Feature('Smarter references');
+        $this->dm->persist($f);
+
+        $q1 = $this->dm->createQueryBuilder($class)
+            ->findAndUpdate()
+            ->field($field)->set($f)
+            ->getQuery()->debug();
+
+        $expected = $this->dm->createDBRef($f, $this->dm->getClassMetadata($class)->fieldMappings[$field]);
+        $this->assertEquals($expected, $q1['newObj']['$set'][$field]);
+    }
+
+    public function provideArrayUpdateOperatorsOnReferenceOne()
+    {
+        yield [ChildA::class, 'featureFull'];
+        yield [ChildB::class, 'featureSimple'];
+        yield [ChildC::class, 'featurePartial'];
+    }
 }
 
 /**