Always prepare full references when querying

src/Persisters/DocumentPersister.php

@@ -41,7 +41,6 @@
 
 use function array_combine;
 use function array_fill;
-use function array_intersect_key;
 use function array_key_exists;
 use function array_keys;
 use function array_map;
@@ -1606,40 +1605,25 @@ private function isInTransaction(array $options): bool
     private function prepareReference(string $fieldName, object $value, array $mapping, bool $inNewObj): array
     {
         $reference = $this->dm->createReference($value, $mapping);
+
+        // If a reference is stored as an identifier, we can always match it
+        // directly. For ReferenceMany, multi-key indexes are used to match
+        // array elements
         if ($inNewObj || $mapping['storeAs'] === ClassMetadata::REFERENCE_STORE_AS_ID) {
             return [[$fieldName, $reference]];
         }
 
-        switch ($mapping['storeAs']) {
-            case ClassMetadata::REFERENCE_STORE_AS_REF:
-                $keys = ['id' => true];
-                break;
-
-            case ClassMetadata::REFERENCE_STORE_AS_DB_REF:
-            case ClassMetadata::REFERENCE_STORE_AS_DB_REF_WITH_DB:
-                $keys = ['$ref' => true, '$id' => true, '$db' => true];
-
-                if ($mapping['storeAs'] === ClassMetadata::REFERENCE_STORE_AS_DB_REF) {
-                    unset($keys['$db']);
-                }
-
-                if (isset($mapping['targetDocument'])) {
-                    unset($keys['$ref'], $keys['$db']);
-                }
-
-                break;
-
-            default:
-                throw new InvalidArgumentException(sprintf('Reference type %s is invalid.', $mapping['storeAs']));
-        }
-
+        // For other ReferenceMany fields, we need to use $elemMatch to find a
+        // single array element that matches all fields
         if ($mapping['type'] === ClassMetadata::MANY) {
-            return [[$fieldName, ['$elemMatch' => array_intersect_key($reference, $keys)]]];
+            return [[$fieldName, ['$elemMatch' => $reference]]];
         }
 
+        // For ReferenceOne fields, we can use multiple conditions on individual
+        // fields, prefixed with the field name of the reference
         return array_map(
             static fn ($key) => [$fieldName . '.' . $key, $reference[$key]],
-            array_keys($keys),
+            array_keys($reference),
         );
     }
 }

tests/Tests/DocumentRepositoryTest.php

@@ -42,7 +42,10 @@ public function testFindByRefOneFull(): void
             ->getUnitOfWork()
             ->getDocumentPersister(User::class)
             ->prepareQueryOrNewObj(['account' => $account]);
-        $expectedQuery = ['account.$id' => new ObjectId($account->getId())];
+        $expectedQuery = [
+            'account.$ref' => $this->dm->getDocumentCollection(Account::class)->getCollectionName(),
+            'account.$id' => new ObjectId($account->getId()),
+        ];
         self::assertEquals($expectedQuery, $query);
 
         self::assertSame($user, $this->dm->getRepository(User::class)->findOneBy(['account' => $account]));
@@ -65,6 +68,7 @@ public function testFindByRefOneWithoutTargetDocumentFull(): void
             'user.$ref' => 'users',
             'user.$id' => new ObjectId($user->getId()),
             'user.$db' => DOCTRINE_MONGODB_DATABASE,
+            'user._doctrine_class_name' => User::class,
         ];
         self::assertEquals($expectedQuery, $query);
 
@@ -87,6 +91,7 @@ public function testFindByRefOneWithoutTargetDocumentStoredAsDbRef(): void
         $expectedQuery = [
             'userDbRef.$ref' => 'users',
             'userDbRef.$id' => new ObjectId($user->getId()),
+            'userDbRef._doctrine_class_name' => User::class,
         ];
         self::assertEquals($expectedQuery, $query);
 
@@ -105,7 +110,15 @@ public function testFindDiscriminatedByRefManyFull(): void
             ->getUnitOfWork()
             ->getDocumentPersister(Developer::class)
             ->prepareQueryOrNewObj(['projects' => $project]);
-        $expectedQuery = ['projects' => ['$elemMatch' => ['$id' => new ObjectId($project->getId())]]];
+        $expectedQuery = [
+            'projects' => [
+                '$elemMatch' => [
+                    '$ref' => $this->dm->getDocumentCollection(SubProject::class)->getCollectionName(),
+                    '$id' => new ObjectId($project->getId()),
+                    'type' => 'sub-project',
+                ],
+            ],
+        ];
         self::assertEquals($expectedQuery, $query);
 
         self::assertSame($developer, $this->dm->getRepository(Developer::class)->findOneBy(['projects' => $project]));
@@ -150,6 +163,7 @@ public function testFindByRefManyFull(): void
         $expectedQuery = [
             'groups' => [
                 '$elemMatch' => [
+                    '$ref' => $this->dm->getDocumentCollection(Group::class)->getCollectionName(),
                     '$id' => new ObjectId($group->getId()),
                 ],
             ],

tests/Tests/Functional/Ticket/GH2825Test.php

@@ -201,6 +201,120 @@ public function testQueryBuilderReplacesReferenceManyCorrectly(): void
         self::assertEquals(['id' => $referenceId], $result['referencedDocumentsStoreAsRef'][0]);
         self::assertEquals(['$ref' => 'GH2825Document', '$id' => $referenceId], $result['referencedDocumentsStoreAsDbRef'][0]);
     }
+
+    public function testQueryBuilderQueriesEmbeddedReferenceOneCorrectlyUsingEquals(): void
+    {
+        $document           = new GH2825Document('document');
+        $document->embedded = new GH2825Embedded('embedded');
+
+        $reference                                 = new GH2825Document('original');
+        $document->embedded->referenceStoreAsId    = $reference;
+        $document->embedded->referenceStoreAsRef   = $reference;
+        $document->embedded->referenceStoreAsDbRef = $reference;
+
+        $this->dm->persist($reference);
+        $this->dm->persist($document);
+
+        $this->dm->flush();
+        $this->dm->clear();
+
+        $reference = $this->dm->find(GH2825Document::class, $reference->id);
+
+        $query = $this->dm->createQueryBuilder(GH2825Document::class)
+            ->field('embedded.referenceStoreAsId')->equals($reference)
+            ->getQuery();
+
+        self::assertEquals(
+            ['embedded.referenceStoreAsId' => new ObjectId($reference->id)],
+            $query->debug('query'),
+        );
+
+        self::assertNotNull($query->getSingleResult());
+
+        $query = $this->dm->createQueryBuilder(GH2825Document::class)
+            ->field('embedded.referenceStoreAsRef')->equals($reference)
+            ->getQuery();
+
+        self::assertEquals(
+            ['embedded.referenceStoreAsRef.id' => new ObjectId($reference->id)],
+            $query->debug('query'),
+        );
+
+        self::assertNotNull($query->getSingleResult());
+
+        $query = $this->dm->createQueryBuilder(GH2825Document::class)
+            ->field('embedded.referenceStoreAsDbRef')->equals($reference)
+            ->getQuery();
+
+        self::assertEquals(
+            [
+                'embedded.referenceStoreAsDbRef.$ref' => $this->dm->getDocumentCollection(GH2825Document::class)->getCollectionName(),
+                'embedded.referenceStoreAsDbRef.$id' => new ObjectId($reference->id),
+            ],
+            $query->debug('query'),
+        );
+
+        self::assertNotNull($query->getSingleResult());
+    }
+
+    public function testQueryBuilderQueriesEmbeddedReferenceManyCorrectlyUsingEquals(): void
+    {
+        $document           = new GH2825Document('document');
+        $document->embedded = new GH2825Embedded('embedded');
+
+        $reference                                             = new GH2825Document('original');
+        $document->embedded->referencedDocumentsStoreAsId[]    = $reference;
+        $document->embedded->referencedDocumentsStoreAsRef[]   = $reference;
+        $document->embedded->referencedDocumentsStoreAsDbRef[] = $reference;
+
+        $this->dm->persist($reference);
+        $this->dm->persist($document);
+
+        $this->dm->flush();
+        $this->dm->clear();
+
+        $reference = $this->dm->find(GH2825Document::class, $reference->id);
+
+        $query = $this->dm->createQueryBuilder(GH2825Document::class)
+            ->field('embedded.referencedDocumentsStoreAsId')->equals($reference)
+            ->getQuery();
+
+        self::assertEquals(
+            ['embedded.referencedDocumentsStoreAsId' => new ObjectId($reference->id)],
+            $query->debug('query'),
+        );
+
+        self::assertNotNull($query->getSingleResult());
+
+        $query = $this->dm->createQueryBuilder(GH2825Document::class)
+            ->field('embedded.referencedDocumentsStoreAsRef')->equals($reference)
+            ->getQuery();
+
+        self::assertEquals(
+            ['embedded.referencedDocumentsStoreAsRef' => ['$elemMatch' => ['id' => new ObjectId($reference->id)]]],
+            $query->debug('query'),
+        );
+
+        self::assertNotNull($query->getSingleResult());
+
+        $query = $this->dm->createQueryBuilder(GH2825Document::class)
+            ->field('embedded.referencedDocumentsStoreAsDbRef')->equals($reference)
+            ->getQuery();
+
+        self::assertEquals(
+            [
+                'embedded.referencedDocumentsStoreAsDbRef' => [
+                    '$elemMatch' => [
+                        '$ref' => $this->dm->getDocumentCollection(GH2825Document::class)->getCollectionName(),
+                        '$id' => new ObjectId($reference->id),
+                    ],
+                ],
+            ],
+            $query->debug('query'),
+        );
+
+        self::assertNotNull($query->getSingleResult());
+    }
 }
 
 #[ODM\Document]

tests/Tests/Persisters/DocumentPersisterGetShardKeyQueryTest.php

@@ -88,7 +88,10 @@ public function testShardByReference(): void
         $method        = new ReflectionMethod($persister, 'getShardKeyQuery');
         $shardKeyQuery = $method->invoke($persister, $o);
 
-        self::assertSame(['reference.$id' => $userId], $shardKeyQuery);
+        self::assertSame([
+            'reference.$ref' => $this->dm->getDocumentCollection(User::class)->getCollectionName(),
+            'reference.$id' => $userId,
+        ], $shardKeyQuery);
     }
 }