Yet another sharding support implementation

doctrine-mongo-mapping.xsd

@@ -43,6 +43,7 @@
       <xs:element name="lifecycle-callbacks" type="odm:lifecycle-callbacks" minOccurs="0"/>
       <xs:element name="also-load-methods" type="odm:also-load-methods" minOccurs="0"/>
       <xs:element name="indexes" type="odm:indexes" minOccurs="0"/>
+      <xs:element name="shard-key" type="odm:shard-key" minOccurs="0"/>
     </xs:sequence>
     <xs:attribute name="db" type="xs:NMTOKEN" />
     <xs:attribute name="name" type="xs:string" />
@@ -60,7 +61,7 @@
   <xs:complexType name="field">
     <xs:sequence>
       <xs:element name="id-generator-option" type="odm:id-generator-option" minOccurs="0" maxOccurs="unbounded" />
-    </xs:sequence>  
+    </xs:sequence>
     <xs:attribute name="id" type="xs:boolean" default="false" />
     <xs:attribute name="name" type="xs:NMTOKEN" />
     <xs:attribute name="type" type="xs:NMTOKEN" default="string" />
@@ -82,7 +83,7 @@
     <xs:attribute name="sparse" type="xs:boolean" />
     <xs:attribute name="unique" type="xs:boolean" />
   </xs:complexType>
-  
+
   <xs:complexType name="id-generator-option">
     <xs:attribute name="name" type="xs:NMTOKEN" use="required"/>
     <xs:attribute name="value" type="xs:string" use="required"/>
@@ -275,6 +276,24 @@
     </xs:sequence>
   </xs:complexType>
 
+  <xs:complexType name="shard-key">
+    <xs:sequence>
+      <xs:element name="key" type="odm:shard-key-key" minOccurs="1" maxOccurs="unbounded"/>
+      <xs:element name="option" type="odm:shard-key-option" minOccurs="0" maxOccurs="unbounded"/>
+    </xs:sequence>
+    <xs:attribute name="unique" type="xs:boolean"/>
+    <xs:attribute name="numInitialChunks" type="xs:integer"/>
+  </xs:complexType>
+
+  <xs:complexType name="shard-key-key">
+    <xs:attribute name="name" type="xs:NMTOKEN" use="required"/>
+    <xs:attribute name="order" type="xs:NMTOKEN" default="asc" />
+  </xs:complexType>
+
+  <xs:complexType name="shard-key-option">
+    <xs:attribute name="name" type="xs:NMTOKEN" use="required"/>
+    <xs:attribute name="value" type="xs:NMTOKEN" use="required" />
+  </xs:complexType>
 
   <xs:complexType name="also-load-method">
     <xs:attribute name="method" type="xs:NMTOKEN" use="required"/>

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

@@ -80,3 +80,4 @@
 require_once __DIR__ . '/PostLoad.php';
 require_once __DIR__ . '/PreFlush.php';
 require_once __DIR__ . '/HasLifecycleCallbacks.php';
+require_once __DIR__ . '/ShardKey.php';

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

@@ -27,5 +27,6 @@ final class Document extends AbstractDocument
     public $repositoryClass;
     public $indexes = array();
     public $requireIndexes = false;
+    public $shardKey;
     public $slaveOkay;
 }

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

@@ -0,0 +1,30 @@
+<?php
+/*
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ * This software consists of voluntary contributions made by many individuals
+ * and is licensed under the MIT license. For more information, see
+ * <http://www.doctrine-project.org>.
+ */
+
+namespace Doctrine\ODM\MongoDB\Mapping\Annotations;
+
+use Doctrine\Common\Annotations\Annotation;
+
+/** @Annotation */
+final class ShardKey extends Annotation
+{
+    public $keys = array();
+    public $unique;
+    public $numInitialChunks;
+}

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

@@ -20,7 +20,6 @@
 namespace Doctrine\ODM\MongoDB\Mapping;
 
 use Doctrine\Instantiator\Instantiator;
-use Doctrine\ODM\MongoDB\LockException;
 
 /**
  * A <tt>ClassMetadata</tt> instance holds all the object-document mapping metadata
@@ -115,6 +114,7 @@ public function __sleep()
             'generatorOptions',
             'idGenerator',
             'indexes',
+            'shardKey',
         );
 
         // The rest of the metadata is only serialized if necessary.

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

@@ -138,6 +138,7 @@ protected function doLoadMetadata($class, $parent, $rootEntityFound, array $nonS
             $this->addInheritedFields($class, $parent);
             $this->addInheritedRelations($class, $parent);
             $this->addInheritedIndexes($class, $parent);
+            $this->setInheritedShardKey($class, $parent);
             $class->setIdentifier($parent->identifier);
             $class->setVersioned($parent->isVersioned);
             $class->setVersionField($parent->versionField);
@@ -336,4 +337,20 @@ private function addInheritedIndexes(ClassMetadata $subClass, ClassMetadata $par
             $subClass->addIndex($index['keys'], $index['options']);
         }
     }
+
+    /**
+     * Adds inherited shard key to the subclass mapping.
+     *
+     * @param ClassMetadata $subClass
+     * @param ClassMetadata $parentClass
+     */
+    private function setInheritedShardKey(ClassMetadata $subClass, ClassMetadata $parentClass)
+    {
+        if ($parentClass->isSharded()) {
+            $subClass->setShardKey(
+                $parentClass->shardKey['keys'],
+                $parentClass->shardKey['options']
+            );
+        }
+    }
 }

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

@@ -19,11 +19,10 @@
 
 namespace Doctrine\ODM\MongoDB\Mapping;
 
-use Doctrine\ODM\MongoDB\Utility\CollectionHelper;
 use Doctrine\ODM\MongoDB\LockException;
-use Doctrine\ODM\MongoDB\Mapping\MappingException;
 use Doctrine\ODM\MongoDB\Proxy\Proxy;
 use Doctrine\ODM\MongoDB\Types\Type;
+use Doctrine\ODM\MongoDB\Utility\CollectionHelper;
 use InvalidArgumentException;
 
 /**
@@ -194,6 +193,11 @@ class ClassMetadataInfo implements \Doctrine\Common\Persistence\Mapping\ClassMet
      */
     public $indexes = array();
 
+    /**
+     * READ-ONLY: Keys and options describing shard key. Only for sharded collections.
+     */
+    public $shardKey;
+
     /**
      * READ-ONLY: Whether or not queries on this document should require indexes.
      */
@@ -518,7 +522,7 @@ public function setCustomRepositoryClass($repositoryClassName)
         if ($this->isEmbeddedDocument) {
             return;
         }
-        
+
         if ($repositoryClassName && strpos($repositoryClassName, '\\') === false && strlen($this->namespace)) {
             $repositoryClassName = $this->namespace . '\\' . $repositoryClassName;
         }
@@ -798,6 +802,67 @@ public function hasIndexes()
         return $this->indexes ? true : false;
     }
 
+    /**
+     * Set shard key for this Document.
+     *
+     * @param array $keys Array of document keys.
+     * @param array $options Array of sharding options.
+     *
+     * @throws MappingException
+     */
+    public function setShardKey(array $keys, array $options = array())
+    {
+        if ($this->inheritanceType == self::INHERITANCE_TYPE_SINGLE_COLLECTION && !is_null($this->shardKey)) {
+            throw MappingException::shardKeyInSingleCollInheritanceSubclass($this->getName());
+        }
+
+        if ($this->isEmbeddedDocument) {
+            throw MappingException::embeddedDocumentCantHaveShardKey($this->getName());
+        }
+
+        foreach ($keys as $field) {
+            if ($this->getTypeOfField($field) == 'increment') {
+                throw MappingException::noIncrementFieldsAllowedInShardKey($this->getName());
+            }
+        }
+
+        $this->shardKey = array(
+            'keys' => array_map(function($value) {
+                if ($value == 1 || $value == -1) {
+                    return (int) $value;
+                }
+                if (is_string($value)) {
+                    $lower = strtolower($value);
+                    if ($lower === 'asc') {
+                        return 1;
+                    } elseif ($lower === 'desc') {
+                        return -1;
+                    }
+                }
+                return $value;
+            }, $keys),
+            'options' => $options
+        );
+    }
+
+    /**
+     * @return array
+     */
+    public function getShardKey()
+    {
+        return $this->shardKey;
+    }
+
+    /**
+     * Checks whether this document has shard key or not.
+     *
+     * @return bool
+     */
+    public function isSharded()
+    {
+        return $this->shardKey ? true : false;
+    }
+
     /**
      * Sets the change tracking policy used by this class.
      *
@@ -1112,7 +1177,7 @@ public function mapField(array $mapping)
         $mapping['isCascadeRefresh'] = in_array('refresh', $cascades);
         $mapping['isCascadeMerge'] = in_array('merge', $cascades);
         $mapping['isCascadeDetach'] = in_array('detach', $cascades);
-        
+
         if (isset($mapping['type']) && $mapping['type'] === 'file') {
             $mapping['file'] = true;
         }
@@ -1160,7 +1225,7 @@ public function mapField(array $mapping)
             && ! empty($mapping['sort']) && ! CollectionHelper::usesSet($mapping['strategy'])) {
             throw MappingException::referenceManySortMustNotBeUsedWithNonSetCollectionStrategy($this->name, $mapping['fieldName'], $mapping['strategy']);
         }
-        
+
         if ($this->isEmbeddedDocument && $mapping['type'] === 'many' && CollectionHelper::isAtomic($mapping['strategy'])) {
             throw MappingException::atomicCollectionStrategyNotAllowed($mapping['strategy'], $this->name, $mapping['fieldName']);
         }
@@ -1514,7 +1579,7 @@ public function setFieldValue($document, $field, $value)
             //so the proxy needs to be loaded first.
             $document->__load();
         }
-        
+
         $this->reflFields[$field]->setValue($document, $value);
     }
 
@@ -1531,7 +1596,7 @@ public function getFieldValue($document, $field)
         if ($document instanceof Proxy && $field !== $this->identifier && ! $document->__isInitialized()) {
             $document->__load();
         }
-        
+
         return $this->reflFields[$field]->getValue($document);
     }
 
@@ -1543,8 +1608,6 @@ public function getFieldValue($document, $field)
      * @return array  The field mapping.
      *
      * @throws MappingException if the $fieldName is not found in the fieldMappings array
-     *
-     * @throws MappingException
      */
     public function getFieldMapping($fieldName)
     {
@@ -1554,6 +1617,26 @@ public function getFieldMapping($fieldName)
         return $this->fieldMappings[$fieldName];
     }
 
+    /**
+     * Gets the field mapping by its DB name.
+     * E.g. it returns identifier's mapping when called with _id.
+     *
+     * @param string $dbFieldName
+     *
+     * @return array
+     * @throws MappingException
+     */
+    public function getFieldMappingByDbFieldName($dbFieldName)
+    {
+        foreach ($this->fieldMappings as $mapping) {
+            if ($mapping['name'] == $dbFieldName) {
+                return $mapping;
+            }
+        }
+
+        throw MappingException::mappingNotFoundByDbName($this->name, $dbFieldName);
+    }
+
     /**
      * Check if the field is not null.
      *
@@ -1701,7 +1784,7 @@ public function isIdGeneratorNone()
      * value to use depending on the column type.
      *
      * @param array $mapping   The version field mapping array
-     * 
+     *
      * @throws LockException
      */
     public function setVersionMapping(array &$mapping)

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

@@ -185,6 +185,11 @@ public function loadMetadataForClass($className, ClassMetadata $class)
             }
         }
 
+        // Set shard key after all fields to ensure we mapped all its keys
+        if (isset($classAnnotations['Doctrine\ODM\MongoDB\Mapping\Annotations\ShardKey'])) {
+            $this->setShardKey($class, $classAnnotations['Doctrine\ODM\MongoDB\Mapping\Annotations\ShardKey']);
+        }
+
         /** @var $method \ReflectionMethod */
         foreach ($reflClass->getMethods(\ReflectionMethod::IS_PUBLIC) as $method) {
             /* Filter for the declaring class only. Callbacks from parent
@@ -240,6 +245,25 @@ private function addIndex(ClassMetadataInfo $class, $index, array $keys = array(
         $class->addIndex($keys, $options);
     }
 
+    /**
+     * @param ClassMetadataInfo $class
+     * @param ODM\ShardKey      $shardKey
+     *
+     * @throws MappingException
+     */
+    private function setShardKey(ClassMetadataInfo $class, ODM\ShardKey $shardKey)
+    {
+        $options = array();
+        $allowed = array('unique', 'numInitialChunks');
+        foreach ($allowed as $name) {
+            if (isset($shardKey->$name)) {
+                $options[$name] = $shardKey->$name;
+            }
+        }
+
+        $class->setShardKey($shardKey->keys, $options);
+    }
+
     /**
      * Factory method for the Annotation Driver
      *