From 3d4917727b5f9c3aa2ccdb55fe5c19668e53b7ba Mon Sep 17 00:00:00 2001
From: Bradley Schofield <ionicisere@gmail.com>
Date: Tue, 3 Dec 2024 15:14:07 +0900
Subject: [PATCH] Address Comments

---
 app/controllers/api/databases.php             | 36 +++++++++++++++++++
 .../e2e/Services/Databases/DatabasesBase.php  |  2 ++
 2 files changed, 38 insertions(+)

diff --git a/app/controllers/api/databases.php b/app/controllers/api/databases.php
index ac8cff8d84..991b774fbb 100644
--- a/app/controllers/api/databases.php
+++ b/app/controllers/api/databases.php
@@ -3781,6 +3781,42 @@ App::patch('/v1/databases/:databaseId/collections/:collectionId/documents')
                 $queries
             )
         );
+        
+        $processDocument = function (Document $collection, Document $document) use (&$processDocument, $dbForProject, $database) {
+            $document->setAttribute('$databaseId', $database->getId());
+            $document->setAttribute('$collectionId', $collection->getId());
+
+            $relationships = \array_filter(
+                $collection->getAttribute('attributes', []),
+                fn ($attribute) => $attribute->getAttribute('type') === Database::VAR_RELATIONSHIP
+            );
+
+            foreach ($relationships as $relationship) {
+                $related = $document->getAttribute($relationship->getAttribute('key'));
+
+                if (empty($related)) {
+                    continue;
+                }
+                if (!\is_array($related)) {
+                    $related = [$related];
+                }
+
+                $relatedCollectionId = $relationship->getAttribute('relatedCollection');
+                $relatedCollection = Authorization::skip(
+                    fn () => $dbForProject->getDocument('database_' . $database->getInternalId(), $relatedCollectionId)
+                );
+
+                foreach ($related as $relation) {
+                    if ($relation instanceof Document) {
+                        $processDocument($relatedCollection, $relation);
+                    }
+                }
+            }
+        };
+
+        foreach ($documents as $document) {
+            $processDocument($collection, $document);
+        }
 
         $response->dynamic(new Document([
             'total' => \count($documents),
diff --git a/tests/e2e/Services/Databases/DatabasesBase.php b/tests/e2e/Services/Databases/DatabasesBase.php
index fa98fd1a15..c262abfab1 100644
--- a/tests/e2e/Services/Databases/DatabasesBase.php
+++ b/tests/e2e/Services/Databases/DatabasesBase.php
@@ -5554,6 +5554,8 @@ trait DatabasesBase
                 Permission::update(Role::user($this->getUser()['$id'])),
                 Permission::delete(Role::user($this->getUser()['$id'])),
             ], $document['$permissions']);
+            $this->assertEquals($collection['body']['$id'], $document['$collectionId']);
+            $this->assertEquals($data['databaseId'], $document['$databaseId']);
         }
 
         // TEST: Check permissions persist