Better status codes and error handling

4 files changed, 23 insertions, 26 deletions

diff --git a/src/collector/db.py b/src/collector/db.py
index 54cdc21..641e5da 100644
--- a/src/collector/db.py
+++ b/src/collector/db.py
@@ -56,7 +56,7 @@ class DBClient:
             print("Could not connect to DB - mongodb://REDACTED_USERNAME:REDACTED_PASSWORD@mongodb:27017/production")
             app_exit(1)
 
-    async def find(self, search_data: SearchInput) -> Optional[List[Dict[str, Any]]]:
+    async def find(self, search_data: SearchInput) -> List[Dict[str, Any]]:
         """Wrap the find() method, handling timeouts and return data type.
 
         :param search_data: Instance of SearchInput.
@@ -67,14 +67,16 @@ class DBClient:
         cursor = self.collection.find(search_data.search)
 
         # Sort on timestamp
+        # TODO: Also sort on IP as well
         cursor.sort("timestamp", -1).limit(search_data.limit).skip(search_data.skip)
 
         try:
             async for document in cursor:
-                data.append(document)
+                if document is not None:
+                    document["_id"] = str(document["_id"])
+                    data.append(document)
 
-            if data:
-                return data
+            return data
 
         except OperationFailure as exc:
             print(f"DB failed to process: {exc.details}")
@@ -87,8 +89,6 @@ class DBClient:
             print(f"DB connection failed: {exc}")
             raise HTTPException(status_code=500, detail="DB connection failed") from exc
 
-        return None
-
     async def find_one(self, object_id: ObjectId) -> Optional[Dict[str, Any]]:
         """Wrap the find_one() method, handling timeouts and return data type.
 
@@ -99,6 +99,7 @@ class DBClient:
         try:
             document = await self.collection.find_one({"_id": object_id})
             if isinstance(document, Dict):
+                document["_id"] = str(document["_id"])
                 return document
             return None
 
@@ -150,15 +151,14 @@ class DBClient:
 
         try:
             result = await self.collection.delete_one({"_id": object_id})
-            if result.deleted_count == 1:
+            if isinstance(result.deleted_count, int) and result.deleted_count == 1:
                 return object_id
+            return None
 
         except BaseException as exc:
             print(f"DB connection failed: {exc}")
             raise HTTPException(status_code=500, detail="DB connection failed") from exc
 
-        return None
-
     async def estimated_document_count(self) -> Optional[int]:
         """Wrap the estimated_document_count() method, handling timeouts and return data type.
 
diff --git a/src/collector/healthcheck.py b/src/collector/healthcheck.py
index 0b9b0e5..3091d98 100644
--- a/src/collector/healthcheck.py
+++ b/src/collector/healthcheck.py
@@ -18,7 +18,7 @@ def check_collector() -> bool:
     req = requests.get(
         "http://localhost:8000/info",
         timeout=3,
-        # verify="./rootCA.crt",
+        # TODO: verify="./rootCA.crt",
     )
 
     if req.status_code != 200:
diff --git a/src/collector/healthcheck.sh b/src/collector/healthcheck.sh
index 3c5e042..aacc906 100755
--- a/src/collector/healthcheck.sh
+++ b/src/collector/healthcheck.sh
@@ -1,11 +1,11 @@
 #!/bin/bash
 
 # If mongodb container
-if [ $1 = "MONGODB" ]
+if [ "$1" = "MONGODB" ]
 then
-  /usr/bin/mongosh --eval 'disableTelemetry()' -u $MONGODB_USERNAME -p $MONGODB_PASSWORD localhost:27017/production /healthcheck-mongodb.js
+  /usr/bin/mongosh --eval 'disableTelemetry()' -u "$MONGODB_USERNAME" -p "$MONGODB_PASSWORD" localhost:27017/production /healthcheck-mongodb.js
   exit $?
 fi
 
 # If collector
-/usr/bin/python3 ./src/collector/healthcheck.py $1 || exit 1
+/usr/bin/python3 ./src/collector/healthcheck.py "$1" || exit 1
diff --git a/src/collector/main.py b/src/collector/main.py
index 72c4ef7..a2d55c7 100755
--- a/src/collector/main.py
+++ b/src/collector/main.py
@@ -48,10 +48,7 @@ async def search(search_data: SearchInput) -> JSONResponse:
     """
     data = await db.find(search_data)
 
-    if data is None:
-        return JSONResponse(content={"status": "error", "message": "Document not found"}, status_code=400)
-
-    return JSONResponse(content={"status": "success", "docs": json_util.dumps(data)})
+    return JSONResponse(content={"status": "success", "docs": data})
 
 
 @app.post("/sc/v0")
@@ -73,7 +70,7 @@ async def create(request: Request) -> JSONResponse:
     key = await db.insert_one(json_data)
 
     if key is None:
-        return JSONResponse(content={"status": "error", "message": "DB error"}, status_code=400)
+        return JSONResponse(content={"status": "error", "message": "DB error"}, status_code=500)
 
     return JSONResponse(content={"status": "success", "key": str(key)})
 
@@ -108,7 +105,7 @@ async def replace(request: Request) -> JSONResponse:  # pylint: disable=too-many
     document = await db.find_one(object_id)
 
     if document is None:
-        return JSONResponse(content={"status": "error", "message": "Document not found"}, status_code=400)
+        return JSONResponse(content={"status": "error", "message": "Document not found"}, status_code=404)
 
     # Ensure valid schema
     del json_data["_id"]
@@ -120,14 +117,14 @@ async def replace(request: Request) -> JSONResponse:  # pylint: disable=too-many
     returned_object_id = await db.replace_one(object_id, json_data)
 
     if returned_object_id is None:
-        return JSONResponse(content={"status": "error", "message": "DB error"}, status_code=400)
+        return JSONResponse(content={"status": "error", "message": "DB error"}, status_code=500)
 
     return JSONResponse(content={"status": "success", "key": str(object_id)})
 
 
 @app.get("/sc/v0/{key}")
 async def get(key: str) -> JSONResponse:
-    """/sc/v0, GET method
+    """/sc/v0/{key}, GET method
 
     :param key: The document key in the database.
     :return: JSONResponse
@@ -136,14 +133,14 @@ async def get(key: str) -> JSONResponse:
     document = await db.find_one(ObjectId(key))
 
     if document is None:
-        return JSONResponse(content={"status": "error", "message": "Document not found"}, status_code=400)
+        return JSONResponse(content={"status": "error", "message": "Document not found"}, status_code=404)
 
-    return JSONResponse(content={"status": "success", "docs": json_util.dumps(document)})
+    return JSONResponse(content={"status": "success", "doc": document})
 
 
 @app.delete("/sc/v0/{key}")
 async def delete(key: str) -> JSONResponse:
-    """/sc/v0, DELETE method
+    """/sc/v0/{key}, DELETE method
 
     :param key: The document key in the database.
     :return: JSONResponse
@@ -151,7 +148,7 @@ async def delete(key: str) -> JSONResponse:
     result = await db.delete_one(ObjectId(key))
 
     if result is None:
-        return JSONResponse(content={"status": "error", "message": "Document not found"}, status_code=400)
+        return JSONResponse(content={"status": "error", "message": "Document not found"}, status_code=404)
 
     return JSONResponse(content={"status": "success", "key": str(key)})
 
@@ -166,6 +163,6 @@ async def info() -> JSONResponse:
     count = await db.estimated_document_count()
 
     if count is None:
-        return JSONResponse(content={"status": "error", "message": "DB error"}, status_code=400)
+        return JSONResponse(content={"status": "error", "message": "DB error"}, status_code=500)
 
     return JSONResponse(content={"status": "success", "Estimated document count": count})