[WIP] Feature sharding

lib/Doctrine/ODM/MongoDB/DocumentManager.php

@@ -534,17 +534,20 @@ public function flush($document = null, array $options = array())
      *
      * @param string $documentName
      * @param string|object $identifier
+     * @param mixed $query
+     * @param array $sort
      * @return mixed|object The document reference.
      */
-    public function getReference($documentName, $identifier)
+    public function getReference($documentName, $identifier, $query = null, array $sort = null)
     {
         $class = $this->metadataFactory->getMetadataFor($documentName);
 
         // Check identity map first, if its already in there just return it.
-        if ($document = $this->unitOfWork->tryGetById($identifier, $class->rootDocumentName)) {
+        if ($identifier && $document = $this->unitOfWork->tryGetById($identifier, $class->rootDocumentName)) {
             return $document;
         }
-        $document = $this->proxyFactory->getProxy($class->name, $identifier);
+
+        $document = $this->proxyFactory->getProxy($class->name, $identifier, $query ?: $identifier, $sort);
         $this->unitOfWork->registerManaged($document, $identifier, array());
 
         return $document;
@@ -677,42 +680,87 @@ public function getClassNameFromDiscriminatorValue(array $mapping, $value)
      *
      * @param mixed $document A document object
      * @param array $referenceMapping Mapping for the field the references the document
+     * @param boolean $forQuery Is the DBRef for a query?
      *
      * @return array A DBRef array
      */
-    public function createDBRef($document, array $referenceMapping = null)
+    public function createDBRef($document, array $referenceMapping = null, $forQuery = false)
     {
         if (!is_object($document)) {
             throw new \InvalidArgumentException('Cannot create a DBRef, the document is not an object');
         }
         $className = get_class($document);
         $class = $this->getClassMetadata($className);
-        $id = $this->unitOfWork->getDocumentIdentifier($document);
+        $id = $class->getDatabaseIdentifierValue($this->unitOfWork->getDocumentIdentifier($document));
 
-        if (isset($referenceMapping['simple']) && $referenceMapping['simple']) {
-            return $class->getDatabaseIdentifierValue($id);
+        if ($referenceMapping && isset($referenceMapping['simple']) && $referenceMapping['simple']) {
+            return $id;
         }
 
-        $dbRef = array(
-            $this->cmd . 'ref' => $class->getCollection(),
-            $this->cmd . 'id'  => $class->getDatabaseIdentifierValue($id),
-            $this->cmd . 'db'  => $this->getDocumentDatabase($className)->getName()
-        );
+        $dbRef = array();
+        if (!$forQuery || null === $referenceMapping || !isset($referenceMapping['targetDocument'])) {
+            $dbRef[$this->cmd . 'ref'] = $class->getCollection();
+            $dbRef[$this->cmd . 'db'] = $this->getDocumentDatabase($className)->getName();
+        }
 
-        if ($class->discriminatorField) {
-            $dbRef[$class->discriminatorField['name']] = $class->discriminatorValue;
+        if ($class->hasQueryFields()) {
+            foreach ($class->queryFields as $field) {
+                $dbRef[$this->cmd . ($class->isIdentifier($field) ? 'id' : $class->getDatabaseFieldName($field))] = $class->getDatabaseFieldValue($class->getFieldValue($document, $field), $field);
+            }
+
+            if (!$forQuery) {
+                $dbRef[$this->cmd . 'id'] = $id;
+            }
+        } else {
+            $dbRef[$this->cmd . 'id'] = $id;
         }
 
-        // add a discriminator value if the referenced document is not mapped explicitely to a targetDocument
-        if ($referenceMapping && ! isset($referenceMapping['targetDocument'])) {
-            $discriminatorField = isset($referenceMapping['discriminatorField']) ? $referenceMapping['discriminatorField'] : '_doctrine_class_name';
-            $discriminatorValue = isset($referenceMapping['discriminatorMap']) ? array_search($class->getName(), $referenceMapping['discriminatorMap']) : $class->getName();
-            $dbRef[$discriminatorField] = $discriminatorValue;
+        if (!$forQuery) {
+            if ($class->discriminatorField) {
+                $dbRef[$class->discriminatorField['name']] = $class->discriminatorValue;
+            }
+
+            // add a discriminator value if the referenced document is not mapped explicitely to a targetDocument
+            if ($referenceMapping && ! isset($referenceMapping['targetDocument'])) {
+                $discriminatorField = isset($referenceMapping['discriminatorField']) ? $referenceMapping['discriminatorField'] : '_doctrine_class_name';
+                $discriminatorValue = isset($referenceMapping['discriminatorMap']) ? array_search($class->getName(), $referenceMapping['discriminatorMap']) : $class->getName();
+                $dbRef[$discriminatorField] = $discriminatorValue;
+            }
         }
 
         return $dbRef;
     }
 
+    /**
+     * Returns a reference query array from a db ref.
+     *
+     * @param mixed $reference
+     * @param string|object $identifier
+     * @param ClassMetadata $metadata
+     * @param array $referenceMapping
+     *
+     * @return array A DBRef query array
+     */
+    public function getReferenceQueryFromDBRef($reference, $identifier, ClassMetadata $metadata, array $referenceMapping)
+    {
+        if (is_scalar($reference) || !$metadata->hasQueryFields() || isset($referenceMapping['simple']) && $referenceMapping['simple']) {
+            $query = $identifier;
+        } else {
+            $query = array();
+            foreach ($metadata->queryFields as $field) {
+                $fieldName = $metadata->isIdentifier($field) ? 'id' : $metadata->getDatabaseFieldName($field);
+                if (isset($reference[$this->cmd . $fieldName])) {
+                    $query[$fieldName] = $metadata->getPHPFieldValue($reference[$this->cmd . $fieldName], $field);
+                }
+            }
+            if (empty($query)) {
+                $query = $identifier;
+            }
+        }
+
+        return $query;
+    }
+
     /**
      * Throws an exception if the DocumentManager is closed or currently not active.
      *

lib/Doctrine/ODM/MongoDB/DocumentNotFoundException.php

@@ -30,8 +30,8 @@
  */
 class DocumentNotFoundException extends MongoDBException
 {
-    public static function documentNotFound($className, $identifier)
+    public static function documentNotFound($className, $query)
     {
-        return new self(sprintf('The "%s" document with identifier "%s" could not be found.', $className, $identifier));
+        return new self(sprintf('The "%s" document with query "%s" could not be found.', $className, json_encode($query)));
     }
 }

lib/Doctrine/ODM/MongoDB/Hydrator/HydratorFactory.php

@@ -238,7 +238,8 @@ private function generateHydratorClass(ClassMetadata $class, $hydratorClassName,
             }
             \$targetMetadata = \$this->dm->getClassMetadata(\$className);
             \$id = \$targetMetadata->getPHPIdentifierValue(\$mongoId);
-            \$return = \$this->dm->getReference(\$className, \$id);
+            \$query = \$this->dm->getReferenceQueryFromDBRef(\$reference, \$id, \$targetMetadata, \$this->class->fieldMappings['%2\$s']);
+            \$return = \$this->dm->getReference(\$className, \$id, \$query);
             \$this->class->reflFields['%2\$s']->setValue(\$document, \$return);
             \$hydratedData['%2\$s'] = \$return;
         }
@@ -252,6 +253,7 @@ private function generateHydratorClass(ClassMetadata $class, $hydratorClassName,
                 if (isset($mapping['repositoryMethod']) && $mapping['repositoryMethod']) {
                     $code .= sprintf(<<<EOF
 
+        /** @ReferenceOne(RepositoryMethod) */
         \$className = \$this->class->fieldMappings['%2\$s']['targetDocument'];
         \$return = \$this->dm->getRepository(\$className)->%3\$s(\$document);
         \$this->class->reflFields['%2\$s']->setValue(\$document, \$return);
@@ -264,19 +266,17 @@ private function generateHydratorClass(ClassMetadata $class, $hydratorClassName,
                     $mapping['repositoryMethod']
                 );
                 } else {
+                    // TODO: Sharding
                     $code .= sprintf(<<<EOF
 
+        /** @ReferenceOne(InverseSide) */
         \$mapping = \$this->class->fieldMappings['%2\$s'];
-        \$className = \$mapping['targetDocument'];
-        \$targetClass = \$this->dm->getClassMetadata(\$mapping['targetDocument']);
-        \$mappedByMapping = \$targetClass->fieldMappings[\$mapping['mappedBy']];
-        \$mappedByFieldName = isset(\$mappedByMapping['simple']) && \$mappedByMapping['simple'] ? \$mapping['mappedBy'] : \$mapping['mappedBy'].'.id';
-        \$criteria = array_merge(
-            array(\$mappedByFieldName => \$data['_id']),
-            isset(\$this->class->fieldMappings['%2\$s']['criteria']) ? \$this->class->fieldMappings['%2\$s']['criteria'] : array() 
-        );
-        \$sort = isset(\$this->class->fieldMappings['%2\$s']['sort']) ? \$this->class->fieldMappings['%2\$s']['sort'] : array();
-        \$return = \$this->unitOfWork->getDocumentPersister(\$className)->load(\$criteria, null, array(), 0, \$sort);
+        \$qb = \$this->dm->createQueryBuilder(\$mapping['targetDocument']);
+        if (isset(\$mapping['criteria'])) {
+            \$qb->setQueryArray(\$mapping['criteria']);
+        }
+        \$qb->field(\$mapping['mappedBy'])->references(\$document);
+        \$return = \$this->dm->getReference(\$mapping['targetDocument'], null, \$qb->getQueryArray(), isset(\$mapping['sort']) ? \$mapping['sort'] : null);
         \$this->class->reflFields['%2\$s']->setValue(\$document, \$return);
         \$hydratedData['%2\$s'] = \$return;
 
@@ -287,9 +287,10 @@ private function generateHydratorClass(ClassMetadata $class, $hydratorClassName,
                     );
                 }
             } elseif ($mapping['association'] === ClassMetadata::REFERENCE_MANY || $mapping['association'] === ClassMetadata::EMBED_MANY) {
+                $comment = $mapping['association'] === ClassMetadata::REFERENCE_MANY ? 'ReferenceMany' : 'EmbedMany';
                 $code .= sprintf(<<<EOF
 
-        /** @Many */
+        /** @%3\$s */
         \$mongoData = isset(\$data['%1\$s']) ? \$data['%1\$s'] : null;
         \$return = new \Doctrine\ODM\MongoDB\PersistentCollection(new \Doctrine\Common\Collections\ArrayCollection(), \$this->dm, \$this->unitOfWork, '$');
         \$return->setHints(\$hints);
@@ -304,7 +305,8 @@ private function generateHydratorClass(ClassMetadata $class, $hydratorClassName,
 EOF
                 ,
                     $mapping['name'],
-                    $mapping['fieldName']
+                    $mapping['fieldName'],
+                    $comment
                 );
             } elseif ($mapping['association'] === ClassMetadata::EMBED_ONE) {
                 $code .= sprintf(<<<EOF

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,43 @@ 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 (boolean) $this->queryFields;
+    }
+
+    /**
+     * Sets the custom query fields for this Document.
+     *
+     * @param array $fields Array of fields for the query.
+     */
+    public function setQueryFields($fields)
+    {
+        $fields = is_array($fields) ? $fields : array($fields);
+        foreach ($fields as $field) {
+            if (strpos($field, '.') !== false) {
+                return new MongoDBException(sprintf('Embedded or reference shard keys not supported. Key "%s" in class "%s"', $field, $this->name));
+            }
+        }
+
+        $this->queryFields = $fields;
+    }
+
     /**
      * Sets the change tracking policy used by this class.
      *
@@ -1123,6 +1165,17 @@ public function addInheritedFieldMapping(array $fieldMapping)
         $this->fieldMappings[$fieldMapping['fieldName']] = $fieldMapping;
     }
 
+    /**
+     * Gets the database field name.
+     *
+     * @param string $fieldName
+     * @return string
+     */
+    public function getDatabaseFieldName($fieldName)
+    {
+        return $this->fieldMappings[$fieldName]['name'];
+    }
+
     /**
      * Checks whether the class has a mapped association with the given field name.
      *
@@ -1289,7 +1342,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);
     }
 
     /**
@@ -1346,6 +1399,30 @@ public function getFieldValue($document, $field)
         return $this->reflFields[$field]->getValue($document);
     }
 
+    /**
+     * Casts the field value to its php type.
+     *
+     * @param mixed $value
+     * @param string $field
+     */
+    public function getPHPFieldValue($value, $field)
+    {
+        $type = $this->fieldMappings[$field]['type'];
+        return Types\Type::getType($type)->convertToPHPValue($value);
+    }
+
+    /**
+     * Casts the field value to its database type.
+     *
+     * @param mixed $value
+     * @param string $field
+     */
+    public function getDatabaseFieldValue($value, $field)
+    {
+        $type = $this->fieldMappings[$field]['type'];
+        return Types\Type::getType($type)->convertToDatabaseValue($value);
+    }
+
     /**
      * Gets the mapping of a field.
      *

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());