From e686f016521826c8764d8851bd594cef94848a6a Mon Sep 17 00:00:00 2001
From: Tom Hacohen <tom@stosb.com>
Date: Sun, 27 Dec 2020 18:34:23 +0200
Subject: [PATCH] Utils: add a utility for getting objects or 404ing.

---
 etebase_fastapi/collection.py |  6 ++----
 etebase_fastapi/utils.py      | 13 +++++++++++++
 2 files changed, 15 insertions(+), 4 deletions(-)
 create mode 100644 etebase_fastapi/utils.py

diff --git a/etebase_fastapi/collection.py b/etebase_fastapi/collection.py
index 229b8c4..2f8d1ed 100644
--- a/etebase_fastapi/collection.py
+++ b/etebase_fastapi/collection.py
@@ -16,6 +16,7 @@ from .authentication import get_authenticated_user
 from .exceptions import ValidationError, transform_validation_error
 from .msgpack import MsgpackRoute, MsgpackResponse
 from .stoken_handler import filter_by_stoken_and_limit, filter_by_stoken, get_stoken_obj, get_queryset_stoken
+from .utils import get_object_or_404
 
 User = get_user_model()
 collection_router = APIRouter(route_class=MsgpackRoute)
@@ -205,10 +206,7 @@ def get_collection_queryset(user: User, queryset: QuerySet) -> QuerySet:
 def get_item_queryset(
     user: User, collection_uid: str, queryset: QuerySet = default_item_queryset
 ) -> t.Tuple[models.Collection, QuerySet]:
-    try:
-        collection = get_collection_queryset(user, models.Collection.objects).get(uid=collection_uid)
-    except models.Collection.DoesNotExist:
-        raise ValidationError("does_not_exist", "Collection does not exist", status_code=status.HTTP_404_NOT_FOUND)
+    collection = get_object_or_404(get_collection_queryset(user, models.Collection.objects), uid=collection_uid)
     # XXX Potentially add this for performance: .prefetch_related('revisions__chunks')
     queryset = queryset.filter(collection__pk=collection.pk, revisions__current=True)
 
diff --git a/etebase_fastapi/utils.py b/etebase_fastapi/utils.py
new file mode 100644
index 0000000..d7f8c09
--- /dev/null
+++ b/etebase_fastapi/utils.py
@@ -0,0 +1,13 @@
+from fastapi import status
+
+from django.db.models import QuerySet
+from django.core.exceptions import ObjectDoesNotExist
+
+from .exceptions import ValidationError
+
+
+def get_object_or_404(queryset: QuerySet, **kwargs):
+    try:
+        return queryset.get(**kwargs)
+    except ObjectDoesNotExist as e:
+        raise ValidationError("does_not_exist", str(e), status_code=status.HTTP_404_NOT_FOUND)