From f12d5cb610ed161b88419615d35745180e712fa1 Mon Sep 17 00:00:00 2001 From: shamoon <4887959+shamoon@users.noreply.github.com> Date: Wed, 11 Feb 2026 23:00:08 -0800 Subject: [PATCH] Unify this a bit --- .../tests/test_api_document_versions.py | 11 ++- src/documents/views.py | 74 +++++++++---------- 2 files changed, 41 insertions(+), 44 deletions(-) diff --git a/src/documents/tests/test_api_document_versions.py b/src/documents/tests/test_api_document_versions.py index 48ed37fa8..8829394a7 100644 --- a/src/documents/tests/test_api_document_versions.py +++ b/src/documents/tests/test_api_document_versions.py @@ -80,20 +80,23 @@ class TestDocumentVersioningApi(DirectoriesMixin, APITestCase): self.assertEqual(resp.status_code, status.HTTP_404_NOT_FOUND) - def test_root_endpoint_returns_404_when_root_document_missing(self) -> None: + def test_root_endpoint_falls_back_when_root_document_missing(self) -> None: doc = Document( title="orphan", checksum="orphan", mime_type="application/pdf", ) - doc.root_document_id = 123 - doc.root_document = None + doc.pk = 123 + doc.root_document_id = 456 + # Simulate a stale FK: id is set but related object is missing. + doc._state.fields_cache["root_document"] = None with mock.patch("documents.views.Document.global_objects") as manager: manager.select_related.return_value.get.return_value = doc resp = self.client.get("/api/documents/123/root/") - self.assertEqual(resp.status_code, status.HTTP_404_NOT_FOUND) + self.assertEqual(resp.status_code, status.HTTP_200_OK) + self.assertEqual(resp.data["root_id"], 123) def test_root_endpoint_returns_403_when_user_lacks_permission(self) -> None: owner = User.objects.create_user(username="owner") diff --git a/src/documents/views.py b/src/documents/views.py index a8237c669..49cf6713d 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -809,6 +809,14 @@ class DocumentViewSet( ) return super().get_serializer(*args, **kwargs) + @staticmethod + def _get_root_doc(doc: Document) -> Document: + # Use root_document_id to avoid a query when this is already a root. + # If root_document isn't available, fall back to the document itself. + if doc.root_document_id is None: + return doc + return doc.root_document or doc + @extend_schema( operation_id="documents_root", responses=inline_serializer( @@ -828,9 +836,7 @@ class DocumentViewSet( except Document.DoesNotExist: raise Http404 - root_doc = doc if doc.root_document_id is None else doc.root_document - if root_doc is None: - raise Http404 + root_doc = self._get_root_doc(doc) if request.user is not None and not has_perms_owner_aware( request.user, "view_document", @@ -899,14 +905,11 @@ class DocumentViewSet( return latest or root_doc def file_response(self, pk, request, disposition): - request_doc = Document.global_objects.select_related("owner").get(id=pk) - root_doc = ( - request_doc - if request_doc.root_document_id is None - else Document.global_objects.select_related("owner").get( - id=request_doc.root_document_id, - ) - ) + request_doc = Document.global_objects.select_related( + "owner", + "root_document", + ).get(id=pk) + root_doc = self._get_root_doc(request_doc) if request.user is not None and not has_perms_owner_aware( request.user, "view_document", @@ -962,14 +965,11 @@ class DocumentViewSet( ) def metadata(self, request, pk=None): try: - request_doc = Document.objects.select_related("owner").get(pk=pk) - root_doc = ( - request_doc - if request_doc.root_document_id is None - else Document.objects.select_related("owner").get( - id=request_doc.root_document_id, - ) - ) + request_doc = Document.objects.select_related( + "owner", + "root_document", + ).get(pk=pk) + root_doc = self._get_root_doc(request_doc) if request.user is not None and not has_perms_owner_aware( request.user, "view_document", @@ -1157,14 +1157,11 @@ class DocumentViewSet( ) def preview(self, request, pk=None): try: - request_doc = Document.objects.select_related("owner").get(id=pk) - root_doc = ( - request_doc - if request_doc.root_document_id is None - else Document.objects.select_related("owner").get( - id=request_doc.root_document_id, - ) - ) + request_doc = Document.objects.select_related( + "owner", + "root_document", + ).get(id=pk) + root_doc = self._get_root_doc(request_doc) if request.user is not None and not has_perms_owner_aware( request.user, "view_document", @@ -1195,14 +1192,11 @@ class DocumentViewSet( @method_decorator(last_modified(thumbnail_last_modified)) def thumb(self, request, pk=None): try: - request_doc = Document.objects.select_related("owner").get(id=pk) - root_doc = ( - request_doc - if request_doc.root_document_id is None - else Document.objects.select_related("owner").get( - id=request_doc.root_document_id, - ) - ) + request_doc = Document.objects.select_related( + "owner", + "root_document", + ).get(id=pk) + root_doc = self._get_root_doc(request_doc) if request.user is not None and not has_perms_owner_aware( request.user, "view_document", @@ -1598,11 +1592,11 @@ class DocumentViewSet( ) def delete_version(self, request, pk=None, version_id=None): try: - root_doc = Document.objects.select_related("owner").get(pk=pk) - if root_doc.root_document_id is not None: - root_doc = Document.objects.select_related("owner").get( - pk=root_doc.root_document_id, - ) + root_doc = Document.objects.select_related( + "owner", + "root_document", + ).get(pk=pk) + root_doc = self._get_root_doc(root_doc) except Document.DoesNotExist: raise Http404