better object id handling

1 files changed, 15 insertions, 1 deletions

diff --git a/src/soc_collector/db.py b/src/soc_collector/db.py
index f537f4a..b10d865 100644
--- a/src/soc_collector/db.py
+++ b/src/soc_collector/db.py
@@ -6,6 +6,8 @@ from dataclasses import dataclass
 
 from fastapi import HTTPException
 from pydantic import BaseModel
+from bson import ObjectId
+from bson.errors import InvalidId
 from pymongo.errors import OperationFailure
 from pymongo import (
     ASCENDING,
@@ -15,7 +17,19 @@ from motor.motor_asyncio import (
     AsyncIOMotorClient,
     AsyncIOMotorCollection,
 )
-from bson import ObjectId
+
+
+def object_id_from_key(key: str) -> ObjectId:
+    """Get ObjectId from key, 400 if invalid ObjectId
+
+    :param key: Key.
+    :return: ObjectId
+    """
+
+    try:
+        return ObjectId(key)
+    except InvalidId as exc:
+        raise HTTPException(status_code=400, detail="Invalid key/object id") from exc
 
 
 class SearchInput(BaseModel):