Schema manager should also know how to handle shard key on reference

doctrine · Mar 9, 2018 · fbebd62 · fbebd62
1 parent 04938ee
commit fbebd62
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 1 deletion.
diff --git a/lib/Doctrine/ODM/MongoDB/SchemaManager.php b/lib/Doctrine/ODM/MongoDB/SchemaManager.php
@@ -6,6 +6,7 @@
 
 use Doctrine\ODM\MongoDB\Mapping\ClassMetadata;
 use Doctrine\ODM\MongoDB\Mapping\ClassMetadataFactory;
+use Doctrine\ODM\MongoDB\Mapping\MappingException;
 use MongoDB\Driver\Exception\RuntimeException;
 use MongoDB\Model\IndexInfo;
 use function array_filter;
@@ -534,10 +535,26 @@ private function runShardCollectionCommand($documentName)
         $shardKey = $class->getShardKey();
         $adminDb = $this->dm->getClient()->selectDatabase('admin');
 
+        $shardKeyPart = [];
+        foreach ($shardKey['keys'] as $key => $order) {
+            if ($class->hasField($key)) {
+                $mapping = $class->getFieldMapping($key);
+                $fieldName = $mapping['name'];
+
+                if ($class->isSingleValuedReference($key)) {
+                    $fieldName = ClassMetadata::getReferenceFieldName($mapping['storeAs'], $fieldName);
+                }
+            } else {
+                $fieldName = $key;
+            }
+
+            $shardKeyPart[$fieldName] = $order;
+        }
+
         $result = $adminDb->command(
             [
                 'shardCollection' => $dbName . '.' . $class->getCollection(),
-                'key'             => $shardKey['keys'],
+                'key'             => $shardKeyPart,
             ]
         )->toArray()[0];
 

diff --git a/tests/Doctrine/ODM/MongoDB/Tests/SchemaManagerTest.php b/tests/Doctrine/ODM/MongoDB/Tests/SchemaManagerTest.php
@@ -555,6 +555,33 @@ public function testEnsureDocumentShardingIgnoresAlreadyShardedError()
         $this->schemaManager->ensureDocumentSharding($classMetadata->getName());
     }
 
+    public function testEnsureDocumentShardingWithShardByReference()
+    {
+        $dbName = DOCTRINE_MONGODB_DATABASE;
+        $classMetadata = $this->dm->getClassMetadata(\Documents\Sharded\ShardedByUser::class);
+        $collectionName = $classMetadata->getCollection();
+        $dbMock = $this->getMockDatabase();
+        $dbMock->method('getName')->willReturn($dbName);
+        $adminDBMock = $this->getMockDatabase();
+        $connMock = $this->getMockConnection();
+        $connMock->method('selectDatabase')->with('admin')->willReturn($adminDBMock);
+        $this->dm->connection = $connMock;
+        $this->dm->documentDatabases = array($classMetadata->getName() => $dbMock);
+
+        $adminDBMock
+            ->expects($this->at(0))
+            ->method('command')
+            ->with(array('enableSharding' => $dbName))
+            ->willReturn(array('ok' => 1));
+        $adminDBMock
+            ->expects($this->at(1))
+            ->method('command')
+            ->with(array('shardCollection' => $dbName . '.' . $collectionName, 'key' => array('db_user.$id' => -1)))
+            ->willReturn(array('ok' => 1));
+
+        $this->schemaManager->ensureDocumentSharding($classMetadata->getName());
+    }
+
     public function testEnableShardingForDb()
     {
         $this->markTestSkipped('Sharding support is still WIP');

diff --git a/tests/Documents/Sharded/ShardedByUser.php b/tests/Documents/Sharded/ShardedByUser.php
@@ -0,0 +1,18 @@
+<?php
+
+namespace Documents\Sharded;
+
+use Doctrine\ODM\MongoDB\Mapping\Annotations as ODM;
+
+/**
+ * @ODM\Document(collection="sharded.users")
+ * @ODM\ShardKey(keys={"user"="desc"})
+ */
+class ShardedByUser
+{
+    /** @ODM\Id */
+    public $id;
+
+    /** @ODM\ReferenceOne(targetDocument="Documents\Sharded\ShardedUser", name="db_user") */
+    public $user;
+}