Handle ObjectID conversion errors

api/database/aircraft.py

@@ -1,10 +1,10 @@
 from typing import Any
 
-from bson import ObjectId
 from fastapi import HTTPException
 from pymongo.errors import WriteError
 
 from database.db import aircraft_collection
+from utils import to_objectid
 from schemas.aircraft import AircraftDisplaySchema, AircraftCreateSchema, aircraft_display_helper, aircraft_add_helper
 
 
@@ -20,7 +20,7 @@ async def retrieve_aircraft(user: str = "") -> list[AircraftDisplaySchema]:
         async for doc in aircraft_collection.find():
             aircraft.append(AircraftDisplaySchema(**aircraft_display_helper(doc)))
     else:
-        async for doc in aircraft_collection.find({"user": ObjectId(user)}):
+        async for doc in aircraft_collection.find({"user": to_objectid(user)}):
             aircraft.append(AircraftDisplaySchema(**aircraft_display_helper(doc)))
 
     return aircraft
@@ -48,7 +48,7 @@ async def retrieve_aircraft_by_id(id: str) -> AircraftDisplaySchema:
     :param tail_no: Tail number of desired aircraft
     :return: Aircraft details
     """
-    aircraft = await aircraft_collection.find_one({"_id": ObjectId(id)})
+    aircraft = await aircraft_collection.find_one({"_id": to_objectid(id)})
 
     if aircraft is None:
         raise HTTPException(404, "Aircraft not found")
@@ -56,7 +56,7 @@ async def retrieve_aircraft_by_id(id: str) -> AircraftDisplaySchema:
     return AircraftDisplaySchema(**aircraft_display_helper(aircraft))
 
 
-async def insert_aircraft(body: AircraftCreateSchema, id: str) -> ObjectId:
+async def insert_aircraft(body: AircraftCreateSchema, id: str) -> to_objectid:
     """
     Insert a new aircraft into the database
 
@@ -77,17 +77,17 @@ async def update_aircraft(body: AircraftCreateSchema, id: str, user: str) -> Air
     :param user: ID of updating user
     :return: Updated aircraft
     """
-    aircraft = await aircraft_collection.find_one({"_id": ObjectId(id)})
+    aircraft = await aircraft_collection.find_one({"_id": to_objectid(id)})
 
     if aircraft is None:
         raise HTTPException(404, "Aircraft not found")
 
-    updated_aircraft = await aircraft_collection.update_one({"_id": ObjectId(id)},
+    updated_aircraft = await aircraft_collection.update_one({"_id": to_objectid(id)},
                                                             {"$set": aircraft_add_helper(body.model_dump(), user)})
     if updated_aircraft is None:
         raise HTTPException(500, "Failed to update aircraft")
 
-    aircraft = await aircraft_collection.find_one({"_id": ObjectId(id)})
+    aircraft = await aircraft_collection.find_one({"_id": to_objectid(id)})
 
     if aircraft is None:
         raise HTTPException(500, "Failed to fetch updated aircraft")
@@ -104,13 +104,13 @@ async def update_aircraft_field(field: str, value: Any, id: str) -> AircraftDisp
     :param id: ID of aircraft to update
     :return: Updated aircraft
     """
-    aircraft = await aircraft_collection.find_one({"_id": ObjectId(id)})
+    aircraft = await aircraft_collection.find_one({"_id": to_objectid(id)})
 
     if aircraft is None:
         raise HTTPException(404, "Aircraft not found")
 
     try:
-        updated_aircraft = await aircraft_collection.update_one({"_id": ObjectId(id)}, {"$set": {field: value}})
+        updated_aircraft = await aircraft_collection.update_one({"_id": to_objectid(id)}, {"$set": {field: value}})
     except WriteError as e:
         raise HTTPException(400, e.details)
 
@@ -127,10 +127,10 @@ async def delete_aircraft(id: str) -> AircraftDisplaySchema:
     :param id: ID of aircraft to delete
     :return: Deleted aircraft information
     """
-    aircraft = await aircraft_collection.find_one({"_id": ObjectId(id)})
+    aircraft = await aircraft_collection.find_one({"_id": to_objectid(id)})
 
     if aircraft is None:
         raise HTTPException(404, "Aircraft not found")
 
-    await aircraft_collection.delete_one({"_id": ObjectId(id)})
+    await aircraft_collection.delete_one({"_id": to_objectid(id)})
     return AircraftDisplaySchema(**aircraft_display_helper(aircraft))

api/database/flights.py

@@ -1,15 +1,15 @@
 import logging
 from datetime import datetime
-from typing import Dict, Union, Any, get_args, List, get_origin, _type_check, get_type_hints
+from typing import Dict, Union
 
 from bson import ObjectId
-from bson.errors import InvalidId
-from fastapi import HTTPException
-from pydantic import parse_obj_as, TypeAdapter, ValidationError, create_model
 
-from schemas.aircraft import AircraftCreateSchema, aircraft_add_helper, AircraftCategory, AircraftClass, \
+from utils import to_objectid
-    aircraft_class_dict, aircraft_category_dict
+from fastapi import HTTPException
-from .aircraft import retrieve_aircraft_by_tail, update_aircraft, update_aircraft_field, retrieve_aircraft
+from pydantic import ValidationError
+
+from schemas.aircraft import aircraft_class_dict, aircraft_category_dict
+from .aircraft import retrieve_aircraft_by_tail, update_aircraft_field
 from .db import flight_collection, aircraft_collection
 from schemas.flight import FlightConciseSchema, FlightDisplaySchema, FlightCreateSchema, flight_display_helper, \
     flight_add_helper, FlightPatchSchema
@@ -34,7 +34,7 @@ async def retrieve_flights(user: str = "", sort: str = "date", order: int = -1,
     """
     filter_options = {}
     if user != "":
-        filter_options["user"] = ObjectId(user)
+        filter_options["user"] = to_objectid(user)
     if filter != "" and filter_val != "":
         if filter not in fs_keys:
             raise HTTPException(400, f"Invalid filter field: {filter}")
@@ -53,7 +53,7 @@ async def retrieve_totals(user: str, start_date: datetime = None, end_date: date
     :param user:
     :return:
     """
-    match: Dict[str, Union[Dict, ObjectId]] = {"user": ObjectId(user)}
+    match: Dict[str, Union[Dict, ObjectId]] = {"user": to_objectid(user)}
 
     if start_date is not None:
         match.setdefault("date", {}).setdefault("$gte", start_date)
@@ -61,7 +61,7 @@ async def retrieve_totals(user: str, start_date: datetime = None, end_date: date
         match.setdefault("date", {}).setdefault("$lte", end_date)
 
     by_class_pipeline = [
-        {"$match": {"user": ObjectId(user)}},
+        {"$match": {"user": to_objectid(user)}},
         {"$lookup": {
             "from": "flight",
             "let": {"aircraft": "$tail_no"},
@@ -104,7 +104,7 @@ async def retrieve_totals(user: str, start_date: datetime = None, end_date: date
     by_class_list = await class_cursor.to_list(None)
 
     totals_pipeline = [
-        {"$match": {"user": ObjectId(user)}},
+        {"$match": {"user": to_objectid(user)}},
         {"$group": {
             "_id": None,
             "time_total": {"$sum": "$time_total"},
@@ -154,7 +154,7 @@ async def retrieve_flight(id: str) -> FlightDisplaySchema:
     :param id: ID of flight to retrieve
     :return: Flight information
     """
-    flight = await flight_collection.find_one({"_id": ObjectId(id)})
+    flight = await flight_collection.find_one({"_id": to_objectid(id)})
 
     if flight is None:
         raise HTTPException(404, "Flight not found")
@@ -193,7 +193,7 @@ async def update_flight(body: FlightCreateSchema, id: str) -> str:
     :param id: ID of flight to update
     :return: ID of updated flight
     """
-    flight = await flight_collection.find_one({"_id": ObjectId(id)})
+    flight = await flight_collection.find_one({"_id": to_objectid(id)})
 
     if flight is None:
         raise HTTPException(404, "Flight not found")
@@ -208,7 +208,7 @@ async def update_flight(body: FlightCreateSchema, id: str) -> str:
         await update_aircraft_field("hobbs", body.hobbs_end, aircraft.id)
 
     # Update flight in database
-    updated_flight = await flight_collection.update_one({"_id": ObjectId(id)}, {"$set": body.model_dump()})
+    updated_flight = await flight_collection.update_one({"_id": to_objectid(id)}, {"$set": body.model_dump()})
 
     if updated_flight is None:
         raise HTTPException(500, "Failed to update flight")
@@ -228,7 +228,7 @@ async def update_flight_fields(id: str, update: dict) -> str:
         if field not in fs_keys:
             raise HTTPException(400, f"Invalid update field: {field}")
 
-    flight = await flight_collection.find_one({"_id": ObjectId(id)})
+    flight = await flight_collection.find_one({"_id": to_objectid(id)})
 
     if flight is None:
         raise HTTPException(404, "Flight not found")
@@ -246,7 +246,7 @@ async def update_flight_fields(id: str, update: dict) -> str:
         if aircraft is None:
             raise HTTPException(404, "Aircraft not found")
 
-    updated_flight = await flight_collection.update_one({"_id": ObjectId(id)}, {"$set": update_dict})
+    updated_flight = await flight_collection.update_one({"_id": to_objectid(id)}, {"$set": update_dict})
 
     if updated_flight is None:
         raise HTTPException(500, "Failed to update flight")
@@ -261,10 +261,10 @@ async def delete_flight(id: str) -> FlightDisplaySchema:
     :param id: ID of flight to delete
     :return: Deleted flight information
     """
-    flight = await flight_collection.find_one({"_id": ObjectId(id)})
+    flight = await flight_collection.find_one({"_id": to_objectid(id)})
 
     if flight is None:
         raise HTTPException(404, "Flight not found")
 
-    await flight_collection.delete_one({"_id": ObjectId(id)})
+    await flight_collection.delete_one({"_id": to_objectid(id)})
     return FlightDisplaySchema(**flight_display_helper(flight))

api/database/img.py

@@ -5,7 +5,7 @@ from gridfs import NoFile
 from .db import db_client as db, files_collection
 
 import motor.motor_asyncio
-from bson import ObjectId
+from utils import to_objectid
 from fastapi import UploadFile, File, HTTPException
 
 fs = motor.motor_asyncio.AsyncIOMotorGridFSBucket(db)
@@ -35,7 +35,7 @@ async def retrieve_image_metadata(image_id: str = "") -> dict:
     :param image_id: ID of image to retrieve metadata of
     :return: Image metadata
     """
-    info = await files_collection.find_one({"_id": ObjectId(image_id)})
+    info = await files_collection.find_one({"_id": to_objectid(image_id)})
 
     if info is None:
         raise HTTPException(404, "Image not found")
@@ -56,7 +56,7 @@ async def retrieve_image(image_id: str = "") -> tuple[io.BytesIO, str]:
 
     stream = io.BytesIO()
     try:
-        await fs.download_to_stream(ObjectId(image_id), stream)
+        await fs.download_to_stream(to_objectid(image_id), stream)
     except NoFile:
         raise HTTPException(404, "Image not found")
 
@@ -73,7 +73,7 @@ async def delete_image(image_id: str = ""):
     :return: True if deleted
     """
     try:
-        await fs.delete(ObjectId(image_id))
+        await fs.delete(to_objectid(image_id))
     except NoFile:
         raise HTTPException(404, "Image not found")
     except Exception as e:

api/database/users.py

@@ -1,6 +1,8 @@
 import logging
 
 from bson import ObjectId
+
+from utils import to_objectid
 from fastapi import HTTPException
 
 from .db import user_collection, flight_collection
@@ -41,7 +43,7 @@ async def get_user_info_id(id: str) -> UserDisplaySchema:
     :param id: ID of user to retrieve
     :return: User information
     """
-    user = await user_collection.find_one({"_id": ObjectId(id)})
+    user = await user_collection.find_one({"_id": to_objectid(id)})
     if user:
         return UserDisplaySchema(**user_helper(user))
 
@@ -65,7 +67,7 @@ async def get_user_system_info_id(id: str) -> UserSystemSchema:
     :param id: ID of user to retrieve
     :return: User information and password
     """
-    user = await user_collection.find_one({"_id": ObjectId(id)})
+    user = await user_collection.find_one({"_id": to_objectid(id)})
     if user:
         return UserSystemSchema(**system_user_helper(user))
 
@@ -89,15 +91,15 @@ async def delete_user(id: str) -> UserDisplaySchema:
     :param id: ID of user to delete
     :return: Information of deleted user
     """
-    user = await user_collection.find_one({"_id": ObjectId(id)})
+    user = await user_collection.find_one({"_id": to_objectid(id)})
 
     if user is None:
         raise HTTPException(404, "User not found")
 
-    await user_collection.delete_one({"_id": ObjectId(id)})
+    await user_collection.delete_one({"_id": to_objectid(id)})
 
     # Delete all flights associated with user
-    await flight_collection.delete_many({"user": ObjectId(id)})
+    await flight_collection.delete_many({"user": to_objectid(id)})
 
     return UserDisplaySchema(**user_helper(user))
 
@@ -127,12 +129,12 @@ async def edit_profile(user_id: str, username: str = None, password: str = None,
             raise HTTPException(403, "Unauthorized attempt to change auth level")
 
     if username:
-        user_collection.update_one({"_id": ObjectId(user_id)}, {"$set": {"username": username}})
+        user_collection.update_one({"_id": to_objectid(user_id)}, {"$set": {"username": username}})
     if password:
         hashed_password = get_hashed_password(password)
-        user_collection.update_one({"_id": ObjectId(user_id)}, {"$set": {"password": hashed_password}})
+        user_collection.update_one({"_id": to_objectid(user_id)}, {"$set": {"password": hashed_password}})
     if auth_level:
-        user_collection.update_one({"_id": ObjectId(user_id)}, {"$set": {"level": auth_level}})
+        user_collection.update_one({"_id": to_objectid(user_id)}, {"$set": {"level": auth_level}})
 
     updated_user = await get_user_info_id(user_id)
     return updated_user

api/database/utils.py

@@ -11,6 +11,7 @@ logger = logging.getLogger("api")
 
 # UTILS #
 
+
 async def create_admin_user():
     """
     Create default admin user if no admin users are present in the database

api/schemas/aircraft.py

@@ -1,6 +1,6 @@
 from enum import Enum
 
-from bson import ObjectId
+from utils import to_objectid
 from pydantic import BaseModel, field_validator
 from pydantic_core.core_schema import ValidationInfo
 
@@ -145,7 +145,7 @@ def aircraft_add_helper(aircraft: dict, user: str) -> dict:
     :param user: User that created aircraft
     :return: Combined dict that can be inserted into db
     """
-    aircraft["user"] = ObjectId(user)
+    aircraft["user"] = to_objectid(user)
     aircraft["aircraft_category"] = aircraft["aircraft_category"].name
     aircraft["aircraft_class"] = aircraft["aircraft_class"].name
 

api/schemas/flight.py

@@ -1,9 +1,7 @@
 import datetime
-import typing
 from typing import Optional, Dict, Union, List
 
-from bson import ObjectId
+from utils import to_objectid
-from fastapi import UploadFile, File
 from pydantic import BaseModel
 
 from schemas.utils import PositiveFloatNullable, PositiveFloat, PositiveInt, PyObjectId
@@ -162,6 +160,6 @@ def flight_add_helper(flight: dict, user: str) -> dict:
     :param user: User that created flight
     :return: Combined dict that can be inserted into db
     """
-    flight["user"] = ObjectId(user)
+    flight["user"] = to_objectid(user)
 
     return flight

api/utils.py

@@ -0,0 +1,17 @@
+from bson import ObjectId
+from bson.errors import InvalidId
+from fastapi import HTTPException
+
+
+def to_objectid(id: str) -> ObjectId:
+    """
+    Try to convert a given string to an ObjectId
+
+    :param id: ID in string form to convert
+    :return: Converted ObjectId
+    """
+    try:
+        oid = ObjectId(id)
+        return oid
+    except InvalidId:
+        raise HTTPException(400, f"{id} is not a recognized ID")