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

doctrine · Feb 15, 2019 · 7c3d409 · 7c3d409
1 parent 06df515
commit 7c3d409
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 1 deletion.
diff --git a/lib/Doctrine/ODM/MongoDB/SchemaManager.php b/lib/Doctrine/ODM/MongoDB/SchemaManager.php
@@ -593,11 +593,27 @@ private function runShardCollectionCommand(string $documentName, ?WriteConcern $
         $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;
+        }
+
         return $adminDb->command(
             array_merge(
                 [
                     'shardCollection' => $dbName . '.' . $class->getCollection(),
-                    'key'             => $shardKey['keys'],
+                    'key'             => $shardKeyPart,
                 ],
                 $this->getWriteOptions(null, $writeConcern)
             )

diff --git a/tests/Doctrine/ODM/MongoDB/Tests/Functional/EnsureShardingTest.php b/tests/Doctrine/ODM/MongoDB/Tests/Functional/EnsureShardingTest.php
@@ -6,6 +6,7 @@
 
 use Doctrine\ODM\MongoDB\MongoDBException;
 use Doctrine\ODM\MongoDB\Tests\BaseTest;
+use Documents\Sharded\ShardedByUser;
 use Documents\Sharded\ShardedOne;
 use Documents\Sharded\ShardedOneWithDifferentKey;
 use function iterator_to_array;
@@ -97,4 +98,21 @@ public function testEnsureShardingForCollectionWithShardingEnabled()
 
         $this->assertTrue($stats['sharded']);
     }
+
+    public function testEnsureDocumentShardingWithShardByReference()
+    {
+        $class = ShardedByUser::class;
+
+        $this->dm->getSchemaManager()->ensureDocumentIndexes($class);
+        $this->dm->getSchemaManager()->ensureDocumentSharding($class);
+
+        $collection = $this->dm->getDocumentCollection($class);
+        $stats      = $this->dm->getDocumentDatabase($class)->command(['collstats' => $collection->getCollectionName()])->toArray()[0];
+        $indexes    = iterator_to_array($collection->listIndexes());
+
+        $this->assertTrue($stats['sharded']);
+
+        $this->assertCount(2, $indexes);
+        $this->assertSame(['db_user.$id' => 1], $indexes[1]->getKey());
+    }
 }
diff --git a/tests/Documents/Sharded/ShardedByUser.php b/tests/Documents/Sharded/ShardedByUser.php
@@ -0,0 +1,20 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Documents\Sharded;
+
+use Doctrine\ODM\MongoDB\Mapping\Annotations as ODM;
+
+/**
+ * @ODM\Document(collection="sharded.users")
+ * @ODM\ShardKey(keys={"user"="asc"})
+ */
+class ShardedByUser
+{
+    /** @ODM\Id */
+    public $id;
+
+    /** @ODM\ReferenceOne(targetDocument="Documents\Sharded\ShardedUser", name="db_user") */
+    public $user;
+}