diff --git a/api/database/aircraft.py b/api/database/aircraft.py index 1619057..092e87b 100644 --- a/api/database/aircraft.py +++ b/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)) diff --git a/api/database/flights.py b/api/database/flights.py index 88840dd..ffa8039 100644 --- a/api/database/flights.py +++ b/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, \ - aircraft_class_dict, aircraft_category_dict -from .aircraft import retrieve_aircraft_by_tail, update_aircraft, update_aircraft_field, retrieve_aircraft +from utils import to_objectid +from fastapi import HTTPException +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)) diff --git a/api/database/img.py b/api/database/img.py index 0eef60f..52e8b13 100644 --- a/api/database/img.py +++ b/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: diff --git a/api/database/users.py b/api/database/users.py index 7d32614..e502ff8 100644 --- a/api/database/users.py +++ b/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 diff --git a/api/database/utils.py b/api/database/utils.py index 51839d4..862fcb5 100644 --- a/api/database/utils.py +++ b/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 diff --git a/api/schemas/aircraft.py b/api/schemas/aircraft.py index 9051b84..ac4ec6e 100644 --- a/api/schemas/aircraft.py +++ b/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 diff --git a/api/schemas/flight.py b/api/schemas/flight.py index f7ce316..fd01315 100644 --- a/api/schemas/flight.py +++ b/api/schemas/flight.py @@ -1,9 +1,7 @@ import datetime -import typing from typing import Optional, Dict, Union, List -from bson import ObjectId -from fastapi import UploadFile, File +from utils import to_objectid 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 diff --git a/api/utils.py b/api/utils.py new file mode 100644 index 0000000..7ac8a5e --- /dev/null +++ b/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")