diff --git a/lib/Doctrine/ODM/MongoDB/SchemaManager.php b/lib/Doctrine/ODM/MongoDB/SchemaManager.php index 3a1633f129..4fe5f0c325 100644 --- 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 index 9f21df6702..5b17cd7759 100644 --- 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 new file mode 100644 index 0000000000..759e5492b7 --- /dev/null +++ b/tests/Documents/Sharded/ShardedByUser.php @@ -0,0 +1,20 @@ +