Update flights to reference aircraft collection

api/database/aircraft.py

@@ -2,8 +2,7 @@ from bson import ObjectId
 from fastapi import HTTPException
 
 from database.db import aircraft_collection
-from database.utils import aircraft_display_helper, aircraft_add_helper
-from schemas.aircraft import AircraftDisplaySchema, AircraftCreateSchema
+from schemas.aircraft import AircraftDisplaySchema, AircraftCreateSchema, aircraft_display_helper, aircraft_add_helper
 
 
 async def retrieve_aircraft(user: str = "") -> list[AircraftDisplaySchema]:

api/database/flights.py

@@ -2,11 +2,12 @@ import logging
 from datetime import datetime
 
 from bson import ObjectId
+from bson.errors import InvalidId
 from fastapi import HTTPException
 
-from database.utils import flight_display_helper, flight_add_helper
 from .db import flight_collection
-from schemas.flight import FlightConciseSchema, FlightDisplaySchema, FlightCreateSchema
+from schemas.flight import FlightConciseSchema, FlightDisplaySchema, FlightCreateSchema, flight_display_helper, \
+    flight_add_helper
 
 logger = logging.getLogger("api")
 
@@ -104,6 +105,14 @@ async def insert_flight(body: FlightCreateSchema, id: str) -> ObjectId:
     :param id: ID of creating user
     :return: ID of inserted flight
     """
+    try:
+        aircraft = await flight_collection.find_one({"_id": ObjectId(body.aircraft)})
+    except InvalidId:
+        raise HTTPException(400, "Invalid aircraft ID")
+
+    if aircraft is None:
+        raise HTTPException(404, "Aircraft not found")
+
     flight = await flight_collection.insert_one(flight_add_helper(body.model_dump(), id))
     return flight.inserted_id
 
@@ -121,6 +130,11 @@ async def update_flight(body: FlightCreateSchema, id: str) -> FlightDisplaySchem
     if flight is None:
         raise HTTPException(404, "Flight not found")
 
+    aircraft = await flight_collection.find_ond({"_id": ObjectId(body.aircraft)})
+
+    if aircraft is None:
+        raise HTTPException(404, "Aircraft not found")
+
     updated_flight = await flight_collection.update_one({"_id": ObjectId(id)}, {"$set": body.model_dump()})
     if updated_flight is None:
         raise HTTPException(500, "Failed to update flight")

api/database/users.py

@@ -3,10 +3,10 @@ import logging
 from bson import ObjectId
 from fastapi import HTTPException
 
-from database.utils import user_helper, create_user_helper, system_user_helper
 from .db import user_collection, flight_collection
 from routes.utils import get_hashed_password
-from schemas.user import UserDisplaySchema, UserCreateSchema, UserSystemSchema, AuthLevel
+from schemas.user import UserDisplaySchema, UserCreateSchema, UserSystemSchema, AuthLevel, user_helper, \
+    create_user_helper, system_user_helper
 
 logger = logging.getLogger("api")
 

api/database/utils.py

@@ -1,114 +1,14 @@
 import logging
 
-from bson import ObjectId
-
 from app.config import get_settings
-from schemas.aircraft import AircraftCategory, AircraftClass
 from .db import user_collection
 from routes.utils import get_hashed_password
 from schemas.user import AuthLevel, UserCreateSchema
+from .users import add_user
 
 logger = logging.getLogger("api")
 
 
-def user_helper(user) -> dict:
-    """
-    Convert given db response into a format usable by UserDisplaySchema
-
-    :param user: Database response
-    :return: Usable dict
-    """
-    return {
-        "id": str(user["_id"]),
-        "username": user["username"],
-        "level": user["level"],
-    }
-
-
-def system_user_helper(user) -> dict:
-    """
-    Convert given db response to a format usable by UserSystemSchema
-
-    :param user: Database response
-    :return: Usable dict
-    """
-    return {
-        "id": str(user["_id"]),
-        "username": user["username"],
-        "password": user["password"],
-        "level": user["level"],
-    }
-
-
-def create_user_helper(user) -> dict:
-    """
-    Convert given db response to a format usable by UserCreateSchema
-
-    :param user: Database response
-    :return: Usable dict
-    """
-    return {
-        "username": user["username"],
-        "password": user["password"],
-        "level": user["level"].value,
-    }
-
-
-def flight_display_helper(flight: dict) -> dict:
-    """
-    Convert given db response to a format usable by FlightDisplaySchema
-
-    :param flight: Database response
-    :return: Usable dict
-    """
-    flight["id"] = str(flight["_id"])
-    flight["user"] = str(flight["user"])
-
-    return flight
-
-
-def flight_add_helper(flight: dict, user: str) -> dict:
-    """
-    Convert given flight schema and user string to a format that can be inserted into the db
-
-    :param flight: Flight request body
-    :param user: User that created flight
-    :return: Combined dict that can be inserted into db
-    """
-    flight["user"] = ObjectId(user)
-    return flight
-
-
-def aircraft_add_helper(aircraft: dict, user: str) -> dict:
-    """
-    Convert given aircraft dict to a format that can be inserted into the db
-
-    :param aircraft: Aircraft request body
-    :param user: User that created aircraft
-    :return: Combined dict that can be inserted into db
-    """
-    aircraft["user"] = ObjectId(user)
-    aircraft["aircraft_category"] = aircraft["aircraft_category"].name
-    aircraft["aircraft_class"] = aircraft["aircraft_class"].name
-    return aircraft
-
-
-def aircraft_display_helper(aircraft: dict) -> dict:
-    """
-    Convert given db response into a format usable by AircraftDisplaySchema
-
-    :param aircraft:
-    :return: USable dict
-    """
-    aircraft["id"] = str(aircraft["_id"])
-    aircraft["user"] = str(aircraft["user"])
-    if aircraft["aircraft_category"] is not AircraftCategory:
-        aircraft["aircraft_category"] = AircraftCategory.__members__.get(aircraft["aircraft_category"])
-    if aircraft["aircraft_class"] is not AircraftClass:
-        aircraft["aircraft_class"] = AircraftClass.__members__.get(aircraft["aircraft_class"])
-    return aircraft
-
-
 # UTILS #
 
 async def create_admin_user():
@@ -131,4 +31,8 @@ async def create_admin_user():
         hashed_password = get_hashed_password(admin_password)
         user = await add_user(
             UserCreateSchema(username=admin_username, password=hashed_password, level=AuthLevel.ADMIN.value))
-        logger.info("Default admin user created with username %s", user.username)
+
+        if user is None:
+            raise Exception("Failed to create default admin user")
+
+        logger.info("Default admin user created with username %s", admin_username)