[WIP] Feature sharding

lib/Doctrine/ODM/MongoDB/Mapping/Annotations/DoctrineAnnotations.php

@@ -61,6 +61,8 @@ final class DiscriminatorField extends Annotation
 final class DiscriminatorMap extends Annotation {}
 /** @Annotation */
 final class DiscriminatorValue extends Annotation {}
+/** @Annotation */
+final class QueryFields extends Annotation {}
 
 /** @Annotation */
 final class Indexes extends Annotation {}

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

@@ -157,6 +157,10 @@ public function __sleep()
             $serialized[] = 'file';
         }
 
+        if($this->hasQueryFields()) {
+            $serialized[] = 'queryFields';
+        }
+
         return $serialized;
     }
 

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

@@ -182,6 +182,11 @@ class ClassMetadataInfo implements \Doctrine\Common\Persistence\Mapping\ClassMet
      */
     public $requireIndexes = false;
 
+    /**
+     * READ-ONLY: The fields for the query for the document collection.
+     */
+    public $queryFields = array();
+
     /**
      * READ-ONLY: The name of the document class.
      */
@@ -681,6 +686,40 @@ public function hasIndexes()
         return $this->indexes ? true : false;
     }
 
+    /**
+     * Returns the custom query fields for this Document.
+     *
+     * @return array $fields The fields for the query.
+     */
+    public function getQueryFields()
+    {
+        return $this->queryFields;
+    }
+
+    /**
+     * Checks whether this document has custom query fields or not.
+     *
+     * @return boolean
+     */
+    public function hasQueryFields()
+    {
+        return $this->queryFields ? true : false;
+    }
+
+    /**
+     * Sets the custom query fields for this Document.
+     *
+     * @param array $fields Array of fields for the query.
+     */
+    public function setQueryFields($fields)
+    {
+        if(!is_array($fields)) {
+            $fields = array($fields);
+        }
+
+        $this->queryFields = $fields;
+    }
+
     /**
      * Sets the change tracking policy used by this class.
      *
@@ -1289,7 +1328,7 @@ public function getIdentifierValue($document)
         if ($document instanceof Proxy && !$document->__isInitialized()) {
             return $document->__identifier__;
         }
-        return (string) $this->reflFields[$this->identifier]->getValue($document);
+        return $this->reflFields[$this->identifier]->getValue($document);
     }
 
     /**

lib/Doctrine/ODM/MongoDB/Mapping/Driver/AnnotationDriver.php

@@ -152,6 +152,8 @@ public function loadMetadataForClass($className, ClassMetadataInfo $class)
                 foreach (is_array($annot->value) ? $annot->value : array($annot->value) as $index) {
                     $this->addIndex($class, $index);
                 }
+            } elseif ($annot instanceof ODM\QueryFields) {
+                $class->setQueryFields($annot->value);
             } elseif ($annot instanceof ODM\InheritanceType) {
                 $class->setInheritanceType(constant('Doctrine\\ODM\\MongoDB\\Mapping\\ClassMetadata::INHERITANCE_TYPE_'.$annot->value));
             } elseif ($annot instanceof ODM\DiscriminatorField) {

lib/Doctrine/ODM/MongoDB/Mapping/Driver/XmlDriver.php

@@ -87,6 +87,9 @@ public function loadMetadataForClass($className, ClassMetadataInfo $class)
             }
             $class->setDiscriminatorMap($map);
         }
+        if (isset($xmlRoot->{'query-fields'})) {
+            $this->setQueryFields($class, $xmlRoot->{'query-fields'});
+        }
         if (isset($xmlRoot->{'indexes'})) {
             foreach($xmlRoot->{'indexes'}->{'index'} as $index) {
                 $this->addIndex($class, $index);
@@ -301,6 +304,15 @@ private function addIndex(ClassMetadataInfo $class, SimpleXmlElement $xmlIndex)
         $class->addIndex($index['keys'], $index['options']);
     }
 
+    private function setQueryFields(ClassMetadataInfo $class, SimpleXmlElement $xmlIndex)
+    {
+        $fields = array();
+        foreach ($xmlIndex->{'key'} as $key) {
+            $fields[(string) $key['name']] = (string) $key['name'];
+        }
+        $class->setQueryFields(array_values($fields));
+    }
+
     protected function loadMappingFile($file)
     {
         $result = array();

lib/Doctrine/ODM/MongoDB/Mapping/Driver/YamlDriver.php

@@ -65,6 +65,9 @@ public function loadMetadataForClass($className, ClassMetadataInfo $class)
         } elseif ($element['type'] === 'embeddedDocument') {
             $class->isEmbeddedDocument = true;
         }
+        if (isset($element['queryFields'])) {
+            $class->setQueryFields($element['queryFields']);
+        }
         if (isset($element['indexes'])) {
             foreach($element['indexes'] as $index) {
                 $class->addIndex($index['keys'], isset($index['options']) ? $index['options'] : array());

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

@@ -240,13 +240,17 @@ public function executeInserts(array $options = array())
             $upsertOptions = $options;
             $upsertOptions['upsert'] = true;
             foreach ($upserts as $oid => $data) {
-                $criteria = array('_id' => $data[$this->cmd.'set']['_id']);
-                unset($data[$this->cmd.'set']['_id']);
+                $query = $this->getDocumentQuery($this->queuedInserts[$oid]);
+                foreach ($query as $field => $value) {
+                    if (isset($data[$this->cmd.'set'][$field])) {
+                        unset($data[$this->cmd.'set'][$field]);
+                    }
+                }
                 // stupid php
                 if (empty($data[$this->cmd.'set'])) {
                     $data[$this->cmd.'set'] = new \stdClass;
                 }
-                $this->collection->update($criteria, $data, $upsertOptions);
+                $this->collection->update($query, $data, $upsertOptions);
             }
         }
 
@@ -263,13 +267,10 @@ public function executeInserts(array $options = array())
      */
     public function update($document, array $options = array())
     {
-        $id = $this->uow->getDocumentIdentifier($document);
         $update = $this->pb->prepareUpdateData($document);
 
         if ( ! empty($update)) {
-
-            $id = $this->class->getDatabaseIdentifierValue($id);
-            $query = array('_id' => $id);
+            $query = $this->getDocumentQuery($document);
 
             // Include versioning logic to set the new version value in the database
             // and to ensure the version has not changed since this document object instance
@@ -320,8 +321,7 @@ public function update($document, array $options = array())
      */
     public function delete($document, array $options = array())
     {
-        $id = $this->uow->getDocumentIdentifier($document);
-        $query = array('_id' => $this->class->getDatabaseIdentifierValue($id));
+        $query = $this->getDocumentQuery($document);
 
         if ($this->class->isLockable) {
             $query[$this->class->lockField] = array($this->cmd . 'exists' => false);
@@ -338,13 +338,12 @@ public function delete($document, array $options = array())
     /**
      * Refreshes a managed document.
      *
-     * @param array $id The identifier of the document.
      * @param object $document The document to refresh.
      */
-    public function refresh($id, $document)
+    public function refresh($document)
     {
-        $class = $this->dm->getClassMetadata(get_class($document));
-        $data = $this->collection->findOne(array('_id' => $id));
+        $query = $this->getDocumentQuery($document);
+        $data = $this->collection->findOne($query);
         $data = $this->hydratorFactory->hydrate($document, $data);
         $this->uow->setOriginalDocumentData($document, $data);
     }
@@ -449,8 +448,8 @@ private function wrapCursor(BaseCursor $cursor)
      */
     public function exists($document)
     {
-        $id = $this->class->getIdentifierObject($document);
-        return (bool) $this->collection->findOne(array(array('_id' => $id)), array('_id'));
+        $query = $this->getDocumentQuery($document);
+        return (bool) $this->collection->findOne($query, array('_id'));
     }
 
     /**
@@ -461,10 +460,9 @@ public function exists($document)
      */
     public function lock($document, $lockMode)
     {
-        $id = $this->uow->getDocumentIdentifier($document);
-        $criteria = array('_id' => $this->class->getDatabaseIdentifierValue($id));
+        $query = $this->getDocumentQuery($document);
         $lockMapping = $this->class->fieldMappings[$this->class->lockField];
-        $this->collection->update($criteria, array($this->cmd.'set' => array($lockMapping['name'] => $lockMode)));
+        $this->collection->update($query, array($this->cmd.'set' => array($lockMapping['name'] => $lockMode)));
         $this->class->reflFields[$this->class->lockField]->setValue($document, $lockMode);
     }
 
@@ -475,10 +473,9 @@ public function lock($document, $lockMode)
      */
     public function unlock($document)
     {
-        $id = $this->uow->getDocumentIdentifier($document);
-        $criteria = array('_id' => $this->class->getDatabaseIdentifierValue($id));
+        $query = $this->getDocumentQuery($document);
         $lockMapping = $this->class->fieldMappings[$this->class->lockField];
-        $this->collection->update($criteria, array($this->cmd.'unset' => array($lockMapping['name'] => true)));
+        $this->collection->update($query, array($this->cmd.'unset' => array($lockMapping['name'] => true)));
         $this->class->reflFields[$this->class->lockField]->setValue($document, null);
     }
 
@@ -740,6 +737,40 @@ private function loadReferenceManyWithRepositoryMethod(PersistentCollection $col
         }
     }
 
+    /**
+     * Gets the query for the document.
+     *
+     * @param mixed $document
+     *
+     * @return array $query
+     */
+    public function getDocumentQuery($document)
+    {
+        if ($this->class->hasQueryFields()) {
+            $changeSet = $this->uow->getDocumentChangeSet($document);
+            $query = array();
+            foreach ($this->class->queryFields as $field) {
+                if (isset($changeSet[$field])) {
+                    if (null !== $changeSet[$field][0]) {
+                        $query[$field] = $changeSet[$field][0];
+                    } else {
+                        // upsert new document
+                        $query[$field] = $changeSet[$field][1];
+                    }
+                } else {
+                    $query[$field] = $this->class->getFieldValue($document, $field);
+                }
+            }
+
+            $query = $this->prepareQuery($query);
+        } else {
+            $id = $this->uow->getDocumentIdentifier($document);
+            $query = array('_id' => $this->class->getDatabaseIdentifierValue($id));
+        }
+
+        return $query;
+    }
+
     /**
      * Prepares a sort array and converts PHP property names to MongoDB field names.
      *

lib/Doctrine/ODM/MongoDB/UnitOfWork.php

@@ -2100,8 +2100,7 @@ private function doRefresh($document, array &$visited)
 
         $class = $this->dm->getClassMetadata(get_class($document));
         if ($this->getDocumentState($document) == self::STATE_MANAGED) {
-            $id = $class->getDatabaseIdentifierValue($this->documentIdentifiers[$oid]);
-            $this->getDocumentPersister($class->name)->refresh($id, $document);
+            $this->getDocumentPersister($class->name)->refresh($document);
         } else {
             throw new \InvalidArgumentException("Document is not MANAGED.");
         }

tests/Doctrine/ODM/MongoDB/Tests/Functional/QueryTest.php

@@ -149,8 +149,12 @@ public function testUpdateQuery()
         $query = $qb->getQuery();
         $result = $query->execute();
 
-        $this->dm->refresh($this->user);
-        $this->assertEquals('crap', $this->user->getUsername());
+        $qb = $this->dm->createQueryBuilder('Documents\User')
+            ->find()
+            ->field('username')->equals('crap');
+        $query = $qb->getQuery();
+        $user = $query->getSingleResult();
+        $this->assertNotNull($user);
     }
 
     public function testUpsertUpdateQuery()

tests/Doctrine/ODM/MongoDB/Tests/Mapping/AbstractMappingDriverTest.php

@@ -217,12 +217,25 @@ public function testIndexes($class)
 
         return $class;
     }
+
+    /**
+     * @depends testCustomFieldName
+     * @param ClassMetadata $class
+     */
+    public function testQueryFields($class)
+    {
+        $this->assertTrue(isset($class->queryFields));
+        $this->assertEquals(array('id', 'name'), $class->queryFields);
+
+        return $class;
+    }
 }
 
 /**
  * @ODM\Document(collection="cms_users")
  * @ODM\DiscriminatorField(fieldName="discr")
  * @ODM\DiscriminatorMap({"default"="Doctrine\ODM\MongoDB\Tests\Mapping\User"})
+ * @ODM\QueryFields({"id", "name"})
  */
 class User
 {
@@ -366,5 +379,6 @@ public static function loadMetadata(ClassMetadata $metadata)
         $metadata->addIndex(array('username' => 'desc'), array('unique' => true));
         $metadata->addIndex(array('email' => 'desc'), array('unique' => true, 'dropDups' => true));
         $metadata->addIndex(array('mysqlProfileId' => 'desc'), array('unique' => true, 'dropDups' => true));
+        $metadata->setQueryFields(array('id', 'name'));
     }
-}
+}

tests/Doctrine/ODM/MongoDB/Tests/Mapping/ClassMetadataTest.php

@@ -29,6 +29,7 @@ public function testClassMetadataInstanceSerialization()
         $cm->setDiscriminatorField(array('name' => 'disc'));
         $cm->mapOneEmbedded(array('fieldName' => 'phonenumbers', 'targetDocument' => 'Bar'));
         $cm->setFile('customFileProperty');
+        $cm->setQueryFields(array('name'));
         $this->assertTrue(is_array($cm->getFieldMapping('phonenumbers')));
         $this->assertEquals(1, count($cm->fieldMappings));
 
@@ -48,6 +49,7 @@ public function testClassMetadataInstanceSerialization()
         $this->assertTrue(is_array($cm->getFieldMapping('phonenumbers')));
         $this->assertEquals(1, count($cm->fieldMappings));
         $this->assertEquals('customFileProperty', $cm->file);
+        $this->assertEquals(array('name'), $cm->queryFields);
         $mapping = $cm->getFieldMapping('phonenumbers');
         $this->assertEquals('Documents\Bar', $mapping['targetDocument']);
     }
@@ -182,4 +184,4 @@ public function testDuplicateFieldAndAssocationMapping2_ThrowsException()
         $this->setExpectedException('Doctrine\ODM\MongoDB\MongoDBException');
         $cm->mapField(array('fieldName' => 'name', 'columnName' => 'name'));
     }
-}
+}

tests/Doctrine/ODM/MongoDB/Tests/Mapping/xml/Doctrine.ODM.MongoDB.Tests.Mapping.User.dcm.xml

@@ -10,6 +10,10 @@
         <discriminator-map>
             <discriminator-mapping value="default" class="Doctrine\ODM\MongoDB\Tests\Mapping\User" />
         </discriminator-map>
+        <query-fields>
+            <key name="id" />
+            <key name="name" />
+        </query-fields>
         <field fieldName="id" id="true" />
         <field fieldName="name" name="username" type="string" />
         <field fieldName="email" type="string" unique="true" drop-dups="true" order="desc" />
@@ -53,4 +57,4 @@
             <lifecycle-callback method="doStuffOnPostPersist" type="postPersist" />
         </lifecycle-callbacks>
     </document>
-</doctrine-mongo-mapping>
+</doctrine-mongo-mapping>

tests/Doctrine/ODM/MongoDB/Tests/Mapping/yaml/Doctrine.ODM.MongoDB.Tests.Mapping.User.dcm.yml

@@ -5,6 +5,7 @@ Doctrine\ODM\MongoDB\Tests\Mapping\User:
     fieldName: discr
   discriminatorMap:
     default: Doctrine\ODM\MongoDB\Tests\Mapping\User
+  queryFields: [id, name]
   fields:
     id:
       type: id
@@ -56,4 +57,4 @@ Doctrine\ODM\MongoDB\Tests\Mapping\User:
 
   lifecycleCallbacks:
     prePersist: [ doStuffOnPrePersist, doOtherStuffOnPrePersistToo ] 
-    postPersist: [ doStuffOnPostPersist ]
+    postPersist: [ doStuffOnPostPersist ]

tests/Documents/User.php

@@ -8,6 +8,7 @@
 /**
  * @ODM\Document(collection="users")
  * @ODM\InheritanceType("COLLECTION_PER_CLASS")
+ * @ODM\QueryFields({"id", "username"})
  */
 class User extends BaseDocument
 {