Unify this a bit

This commit is contained in:
shamoon
2026-02-11 23:00:08 -08:00
parent 74ce218b78
commit f12d5cb610
2 changed files with 41 additions and 44 deletions

View File

@@ -80,20 +80,23 @@ class TestDocumentVersioningApi(DirectoriesMixin, APITestCase):
self.assertEqual(resp.status_code, status.HTTP_404_NOT_FOUND) 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( doc = Document(
title="orphan", title="orphan",
checksum="orphan", checksum="orphan",
mime_type="application/pdf", mime_type="application/pdf",
) )
doc.root_document_id = 123 doc.pk = 123
doc.root_document = None 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: with mock.patch("documents.views.Document.global_objects") as manager:
manager.select_related.return_value.get.return_value = doc manager.select_related.return_value.get.return_value = doc
resp = self.client.get("/api/documents/123/root/") 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: def test_root_endpoint_returns_403_when_user_lacks_permission(self) -> None:
owner = User.objects.create_user(username="owner") owner = User.objects.create_user(username="owner")

View File

@@ -809,6 +809,14 @@ class DocumentViewSet(
) )
return super().get_serializer(*args, **kwargs) 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( @extend_schema(
operation_id="documents_root", operation_id="documents_root",
responses=inline_serializer( responses=inline_serializer(
@@ -828,9 +836,7 @@ class DocumentViewSet(
except Document.DoesNotExist: except Document.DoesNotExist:
raise Http404 raise Http404
root_doc = doc if doc.root_document_id is None else doc.root_document root_doc = self._get_root_doc(doc)
if root_doc is None:
raise Http404
if request.user is not None and not has_perms_owner_aware( if request.user is not None and not has_perms_owner_aware(
request.user, request.user,
"view_document", "view_document",
@@ -899,14 +905,11 @@ class DocumentViewSet(
return latest or root_doc return latest or root_doc
def file_response(self, pk, request, disposition): def file_response(self, pk, request, disposition):
request_doc = Document.global_objects.select_related("owner").get(id=pk) request_doc = Document.global_objects.select_related(
root_doc = ( "owner",
request_doc "root_document",
if request_doc.root_document_id is None ).get(id=pk)
else Document.global_objects.select_related("owner").get( root_doc = self._get_root_doc(request_doc)
id=request_doc.root_document_id,
)
)
if request.user is not None and not has_perms_owner_aware( if request.user is not None and not has_perms_owner_aware(
request.user, request.user,
"view_document", "view_document",
@@ -962,14 +965,11 @@ class DocumentViewSet(
) )
def metadata(self, request, pk=None): def metadata(self, request, pk=None):
try: try:
request_doc = Document.objects.select_related("owner").get(pk=pk) request_doc = Document.objects.select_related(
root_doc = ( "owner",
request_doc "root_document",
if request_doc.root_document_id is None ).get(pk=pk)
else Document.objects.select_related("owner").get( root_doc = self._get_root_doc(request_doc)
id=request_doc.root_document_id,
)
)
if request.user is not None and not has_perms_owner_aware( if request.user is not None and not has_perms_owner_aware(
request.user, request.user,
"view_document", "view_document",
@@ -1157,14 +1157,11 @@ class DocumentViewSet(
) )
def preview(self, request, pk=None): def preview(self, request, pk=None):
try: try:
request_doc = Document.objects.select_related("owner").get(id=pk) request_doc = Document.objects.select_related(
root_doc = ( "owner",
request_doc "root_document",
if request_doc.root_document_id is None ).get(id=pk)
else Document.objects.select_related("owner").get( root_doc = self._get_root_doc(request_doc)
id=request_doc.root_document_id,
)
)
if request.user is not None and not has_perms_owner_aware( if request.user is not None and not has_perms_owner_aware(
request.user, request.user,
"view_document", "view_document",
@@ -1195,14 +1192,11 @@ class DocumentViewSet(
@method_decorator(last_modified(thumbnail_last_modified)) @method_decorator(last_modified(thumbnail_last_modified))
def thumb(self, request, pk=None): def thumb(self, request, pk=None):
try: try:
request_doc = Document.objects.select_related("owner").get(id=pk) request_doc = Document.objects.select_related(
root_doc = ( "owner",
request_doc "root_document",
if request_doc.root_document_id is None ).get(id=pk)
else Document.objects.select_related("owner").get( root_doc = self._get_root_doc(request_doc)
id=request_doc.root_document_id,
)
)
if request.user is not None and not has_perms_owner_aware( if request.user is not None and not has_perms_owner_aware(
request.user, request.user,
"view_document", "view_document",
@@ -1598,11 +1592,11 @@ class DocumentViewSet(
) )
def delete_version(self, request, pk=None, version_id=None): def delete_version(self, request, pk=None, version_id=None):
try: try:
root_doc = Document.objects.select_related("owner").get(pk=pk) root_doc = Document.objects.select_related(
if root_doc.root_document_id is not None: "owner",
root_doc = Document.objects.select_related("owner").get( "root_document",
pk=root_doc.root_document_id, ).get(pk=pk)
) root_doc = self._get_root_doc(root_doc)
except Document.DoesNotExist: except Document.DoesNotExist:
raise Http404 raise Http404