diff --git a/api/database/img.py b/api/database/img.py index 52e8b13..42d32d8 100644 --- a/api/database/img.py +++ b/api/database/img.py @@ -1,4 +1,6 @@ import io +import mimetypes +import os from gridfs import NoFile @@ -40,20 +42,21 @@ async def retrieve_image_metadata(image_id: str = "") -> dict: if info is None: raise HTTPException(404, "Image not found") - return info["metadata"] + file_extension = os.path.splitext(info["filename"])[1] + media_type = "image/webp" if file_extension == ".webp" else mimetypes.types_map.get(file_extension) + + return {**info["metadata"], 'contentType': media_type if media_type else ""} -async def retrieve_image(image_id: str = "") -> tuple[io.BytesIO, str]: +async def retrieve_image(image_id: str = "") -> tuple[io.BytesIO, str, 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 + :return: BytesIO stream of image file, ID of user that uploaded the image, file type """ metadata = await retrieve_image_metadata(image_id) - print(metadata) - stream = io.BytesIO() try: await fs.download_to_stream(to_objectid(image_id), stream) @@ -62,7 +65,7 @@ async def retrieve_image(image_id: str = "") -> tuple[io.BytesIO, str]: stream.seek(0) - return stream, metadata["user"] if metadata["user"] else "" + return stream, metadata["user"] if metadata["user"] else "", metadata["contentType"] async def delete_image(image_id: str = ""): diff --git a/api/routes/flights.py b/api/routes/flights.py index 499faa4..330c39e 100644 --- a/api/routes/flights.py +++ b/api/routes/flights.py @@ -1,6 +1,6 @@ import logging from datetime import datetime -from typing import Any +from typing import Any, List from fastapi import APIRouter, HTTPException, Depends, Form, UploadFile, File @@ -130,8 +130,8 @@ async def add_flight(flight_body: FlightSchema, user: UserDisplaySchema = Depend return {"id": str(flight)} -@router.post('/{flight_id}/add_images', summary="Add images to a flight log") -async def add_images(log_id: str, images: list[UploadFile] = File(...), +@router.post('/{log_id}/add_images', summary="Add images to a flight log") +async def add_images(log_id: str, images: List[UploadFile] = File(...), user: UserDisplaySchema = Depends(get_current_user)): """ Add images to a flight logbook entry diff --git a/api/routes/img.py b/api/routes/img.py index a0287e8..1df3f4d 100644 --- a/api/routes/img.py +++ b/api/routes/img.py @@ -24,15 +24,12 @@ async def get_image(user: UserDisplaySchema = Depends(get_current_user), :param image_id: ID of image to retrieve :return: Stream associated with requested image """ - stream, user_created = await img.retrieve_image(image_id) + stream, user_created, media_type = await img.retrieve_image(image_id) if not user.id == user_created and not user.level == AuthLevel.ADMIN: raise HTTPException(403, "Access denied") - file_extension = os.path.splitext(image_id)[1] - media_type = mimetypes.types_map.get(file_extension) - - return StreamingResponse(stream, media_type=media_type) + return StreamingResponse(stream, headers={'Content-Type': media_type}) @router.post("/upload", description="Upload an image to the database")