From a1aae6658efa8f22e92b9f3bf52ae49a373cf62b Mon Sep 17 00:00:00 2001
From: Andreas Braun <git@alcaeus.org>
Date: Fri, 10 May 2024 11:34:51 +0200
Subject: [PATCH] Preserve collection changes until transaction has succeeded

---
 .../MongoDB/Persisters/DocumentPersister.php  |  9 ++----
 lib/Doctrine/ODM/MongoDB/UnitOfWork.php       |  6 ++++
 ...WorkTransactionalCommitConsistencyTest.php | 31 +++++++++++++++++++
 3 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/lib/Doctrine/ODM/MongoDB/Persisters/DocumentPersister.php b/lib/Doctrine/ODM/MongoDB/Persisters/DocumentPersister.php
index 2a3a76731..c2adf2ddd 100644
--- a/lib/Doctrine/ODM/MongoDB/Persisters/DocumentPersister.php
+++ b/lib/Doctrine/ODM/MongoDB/Persisters/DocumentPersister.php
@@ -1512,14 +1512,11 @@ private function handleCollections(object $document, array $options): void
             $collections[] = $coll;
         }
 
-        if (! empty($collections)) {
-            $this->cp->update($document, $collections, $options);
+        if (empty($collections)) {
+            return;
         }
 
-        // Take new snapshots from visited collections
-        foreach ($this->uow->getVisitedCollections($document) as $coll) {
-            $coll->takeSnapshot();
-        }
+        $this->cp->update($document, $collections, $options);
     }
 
     /**
diff --git a/lib/Doctrine/ODM/MongoDB/UnitOfWork.php b/lib/Doctrine/ODM/MongoDB/UnitOfWork.php
index c80b9d909..c3b2af5e4 100644
--- a/lib/Doctrine/ODM/MongoDB/UnitOfWork.php
+++ b/lib/Doctrine/ODM/MongoDB/UnitOfWork.php
@@ -477,6 +477,12 @@ function (Session $session) use ($options): void {
             $this->evm->dispatchEvent(Events::postFlush, new Event\PostFlushEventArgs($this->dm));
 
             // Clear up
+            foreach ($this->visitedCollections as $collections) {
+                foreach ($collections as $coll) {
+                    $coll->takeSnapshot();
+                }
+            }
+
             $this->scheduledDocumentInsertions  =
             $this->scheduledDocumentUpserts     =
             $this->scheduledDocumentUpdates     =
diff --git a/tests/Doctrine/ODM/MongoDB/Tests/UnitOfWorkTransactionalCommitConsistencyTest.php b/tests/Doctrine/ODM/MongoDB/Tests/UnitOfWorkTransactionalCommitConsistencyTest.php
index 6a3a3ae27..93df2c431 100644
--- a/tests/Doctrine/ODM/MongoDB/Tests/UnitOfWorkTransactionalCommitConsistencyTest.php
+++ b/tests/Doctrine/ODM/MongoDB/Tests/UnitOfWorkTransactionalCommitConsistencyTest.php
@@ -9,6 +9,7 @@
 use Documents\Address;
 use Documents\ForumUser;
 use Documents\FriendUser;
+use Documents\Phonenumber;
 use Documents\User;
 use MongoDB\BSON\ObjectId;
 use MongoDB\Client;
@@ -561,6 +562,36 @@ public function testTransientDeleteErrorWithEmbeddedDocument(): void
         self::assertFalse($this->uow->isScheduledForDelete($user));
     }
 
+    public function testTransientErrorPreservesCollectionChangesets(): void
+    {
+        // Create a dummy user so that we later have a user to remove
+        $fooUser = new User();
+        $fooUser->setUsername('foo');
+        $this->uow->persist($fooUser);
+        $this->uow->commit();
+
+        // Create a new user with a collection update
+        $alcaeus = new User();
+        // Set an identifier to force an upsert, which causes separate queries
+        // for collection updates
+        $alcaeus->setId(new ObjectId());
+        $alcaeus->setUsername('alcaeus');
+        $alcaeus->getPhonenumbers()->add(new Phonenumber('12345'));
+        $this->uow->persist($alcaeus);
+
+        // Remove fooUser and create a transient failpoint to force the deletion
+        // to fail. This exposes the issue with collections
+        $this->uow->remove($fooUser);
+        $this->createTransientFailPoint('delete');
+
+        $this->uow->commit();
+
+        $this->dm->clear();
+        $check = $this->dm->find(User::class, $alcaeus->getId());
+        self::assertNotNull($check);
+        self::assertCount(1, $check->getPhonenumbers());
+    }
+
     /** Create a document manager with a single host to ensure failpoints target the correct server */
     protected static function createTestDocumentManager(): DocumentManager
     {