Implement image uploads and flight patching

api/database/db.py

@@ -25,4 +25,5 @@ except Exception as e:
 user_collection = db_client["user"]
 flight_collection = db_client["flight"]
 aircraft_collection = db_client["aircraft"]
+files_collection = db_client.fs.files
 token_collection = db_client["token_blacklist"]

api/database/flights.py

@@ -1,20 +1,24 @@
 import logging
 from datetime import datetime
-from typing import Dict, Union
+from typing import Dict, Union, Any, get_args, List, get_origin, _type_check, get_type_hints
 
 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, \
     aircraft_class_dict, aircraft_category_dict
 from .aircraft import retrieve_aircraft_by_tail, update_aircraft, update_aircraft_field, retrieve_aircraft
 from .db import flight_collection, aircraft_collection
 from schemas.flight import FlightConciseSchema, FlightDisplaySchema, FlightCreateSchema, flight_display_helper, \
-    flight_add_helper
+    flight_add_helper, FlightPatchSchema
 
 logger = logging.getLogger("api")
 
+fs_keys = list(FlightCreateSchema.__annotations__.keys())
+fs_keys.extend(list(FlightDisplaySchema.__annotations__.keys()))
+
 
 async def retrieve_flights(user: str = "", sort: str = "date", order: int = -1, filter: str = "",
                            filter_val: str = "") -> list[FlightConciseSchema]:
@@ -32,8 +36,6 @@ async def retrieve_flights(user: str = "", sort: str = "date", order: int = -1,
     if user != "":
         filter_options["user"] = ObjectId(user)
     if filter != "" and filter_val != "":
-        fs_keys = list(FlightCreateSchema.__annotations__.keys())
-        fs_keys.extend(list(FlightDisplaySchema.__annotations__.keys()))
         if filter not in fs_keys:
             raise HTTPException(400, f"Invalid filter field: {filter}")
         filter_options[filter] = filter_val
@@ -214,6 +216,44 @@ async def update_flight(body: FlightCreateSchema, id: str) -> str:
     return id
 
 
+async def update_flight_fields(id: str, update: dict) -> str:
+    """
+    Update a single field of the given flight in the database
+
+    :param id: ID of flight to update
+    :param update: Dictionary of fields and values to update
+    :return: ID of updated flight
+    """
+    for field in update.keys():
+        if field not in fs_keys:
+            raise HTTPException(400, f"Invalid update field: {field}")
+
+    flight = await flight_collection.find_one({"_id": ObjectId(id)})
+
+    if flight is None:
+        raise HTTPException(404, "Flight not found")
+
+    try:
+        parsed_update = FlightPatchSchema.model_validate(update)
+    except ValidationError as e:
+        raise HTTPException(422, e.errors())
+
+    update_dict = {field: value for field, value in parsed_update.model_dump().items() if field in update.keys()}
+
+    if "aircraft" in update_dict.keys():
+        aircraft = await retrieve_aircraft_by_tail(update_dict["aircraft"])
+
+        if aircraft is None:
+            raise HTTPException(404, "Aircraft not found")
+
+    updated_flight = await flight_collection.update_one({"_id": ObjectId(id)}, {"$set": update_dict})
+
+    if updated_flight is None:
+        raise HTTPException(500, "Failed to update flight")
+
+    return id
+
+
 async def delete_flight(id: str) -> FlightDisplaySchema:
     """
     Delete the given flight from the database

api/database/img.py

@@ -0,0 +1,82 @@
+import io
+
+from gridfs import NoFile
+
+from .db import db_client as db, files_collection
+
+import motor.motor_asyncio
+from bson import ObjectId
+from fastapi import UploadFile, File, HTTPException
+
+fs = motor.motor_asyncio.AsyncIOMotorGridFSBucket(db)
+
+
+async def upload_image(image: UploadFile = File(...), user: str = "") -> dict:
+    """
+    Take an image file and add it to the database, returning the filename and ID of the added image
+
+    :param image: Image to upload
+    :param user: ID of user uploading image to encode in image metadata
+    :return: Dictionary with filename and file_id of newly added image
+    """
+    image_data = await image.read()
+
+    metadata = {"user": user}
+
+    file_id = await fs.upload_from_stream(image.filename, io.BytesIO(image_data), metadata=metadata)
+
+    return {"filename": image.filename, "file_id": str(file_id)}
+
+
+async def retrieve_image_metadata(image_id: str = "") -> dict:
+    """
+    Retrieve the metadata of a given image
+
+    :param image_id: ID of image to retrieve metadata of
+    :return: Image metadata
+    """
+    info = await files_collection.find_one({"_id": ObjectId(image_id)})
+
+    if info is None:
+        raise HTTPException(404, "Image not found")
+
+    return info["metadata"]
+
+
+async def retrieve_image(image_id: str = "") -> tuple[io.BytesIO, str]:
+    """
+    Retrieve the given image file from the database along with the user who created it
+
+    :param image_id: ID of image to retrieve
+    :return: BytesIO stream of image file, ID of user that uploaded the image
+    """
+    metadata = await retrieve_image_metadata(image_id)
+
+    print(metadata)
+
+    stream = io.BytesIO()
+    try:
+        await fs.download_to_stream(ObjectId(image_id), stream)
+    except NoFile:
+        raise HTTPException(404, "Image not found")
+
+    stream.seek(0)
+
+    return stream, metadata["user"] if metadata["user"] else ""
+
+
+async def delete_image(image_id: str = ""):
+    """
+    Delete the given image from the database
+
+    :param image_id: ID of image to delete
+    :return: True if deleted
+    """
+    try:
+        await fs.delete(ObjectId(image_id))
+    except NoFile:
+        raise HTTPException(404, "Image not found")
+    except Exception as e:
+        raise HTTPException(500, e)
+
+    return True