from datetime import timedelta from typing import List, Optional, Annotated, TYPE_CHECKING import strawberry from sqlalchemy import select, delete from sqlalchemy.ext.asyncio import AsyncSession from database import models from graphql_schema.dataloaders import copilots_dataloader from graphql_schema.dataloaders.aircraft import aircraft_dataloader from graphql_schema.dataloaders.airport import airport_dataloader from graphql_schema.dataloaders.photos import photos_dataloader, cover_photo_loader from graphql_schema.dataloaders.poi import flight_track_dataloader, poi_dataloader from graphql_schema.entities.aircraft import Aircraft from graphql_schema.entities.airport import Airport from graphql_schema.entities.copilot import CopilotType from graphql_schema.entities.photo import Photo from graphql_schema.entities.poi import PointOfInterest from graphql_schema.sqlalchemy_to_strawberry_type import strawberry_sqlalchemy_type, strawberry_sqlalchemy_input if TYPE_CHECKING: from .copilot import CopilotType @strawberry.input() class PointOfInterestInput: id: Optional[int] = None name: str @strawberry.input() class CopilotInput: id: Optional[int] = None name: str @strawberry_sqlalchemy_type(models.FlightTrack) class FlightTrack: async def load_poi(root): return await poi_dataloader.load(root.point_of_interest_id) point_of_interest: PointOfInterest = strawberry.field(resolver=load_poi) @strawberry_sqlalchemy_type(models.Flight) class Flight: async def load_takeoff_airport(root): return await airport_dataloader.load(root.takeoff_airport_id) async def load_track(root): return await flight_track_dataloader.load(root.id) async def load_landing_airport(root): return await airport_dataloader.load(root.landing_airport_id) async def load_aircraft(root): return await aircraft_dataloader.load(root.aircraft_id) async def load_copilot(root): return await copilots_dataloader.load(root.copilot_id) async def load_photos(root): return await photos_dataloader.load(root.id) async def load_cover_photo(root): return await cover_photo_loader.load(root.id) def duration_min_calculated(root): if root.duration_total: return root.duration_total if root.takeoff_datetime and root.landing_datetime: diff: timedelta = root.landing_datetime - root.takeoff_datetime return diff.seconds / 60 return 0 duration_min_calculated: int = strawberry.field(resolver=duration_min_calculated) copilot: Optional[Annotated["CopilotType", strawberry.lazy(".copilot")]] = strawberry.field(resolver=load_copilot) aircraft: Aircraft = strawberry.field(resolver=load_aircraft) takeoff_airport: Airport = strawberry.field(resolver=load_takeoff_airport) landing_airport: Airport = strawberry.field(resolver=load_landing_airport) cover_photo: Optional[Photo] = strawberry.field(resolver=load_cover_photo) track: List[FlightTrack] = strawberry.field(resolver=load_track) photos: List[Photo] = strawberry.field(resolver=load_photos) def get_base_query(user_id: int): return ( select(models.Flight) .filter(models.Flight.created_by_id == user_id) .filter(models.Flight.deleted.is_(False)) .order_by(models.Flight.id.desc()) ) @strawberry.type class FlightQueries: @strawberry.input class FlightFilters: takeoff: Optional[int] @strawberry.field async def flights(root, info, filters: Optional[FlightFilters] = None) -> List[Flight]: query = get_base_query(info.context.user_id).order_by(models.Flight.id.desc()) return (await info.context.db.scalars(query)).all() @strawberry.field async def flight(root, info, id: int) -> Flight: query = ( get_base_query(info.context.user_id) .filter(models.Flight.id == id) .filter(models.Flight.deleted.is_(False)) ) return (await info.context.db.scalars(query)).one() @strawberry.type class CreateFlightMutation: @strawberry_sqlalchemy_input(models.Flight, exclude_fields=["id"]) class CreateFlightInput: pass @strawberry.mutation async def create_flight(self, info, input: CreateFlightInput) -> Flight: return await models.Flight.create(info.context.db, data={ **input.to_dict(), "created_by_id": info.context.user_id }) async def handle_track_edit(db: AsyncSession, flight: models.Flight, track: List[PointOfInterestInput], user_id: int): await db.execute(delete(models.FlightTrack).filter(models.FlightTrack.flight_id == flight.id)) existing_poi_ids = [i.id for i in track if i.id] poi_query = ( select(models.PointOfInterest) .filter(models.PointOfInterest.created_by_id == user_id) .filter(models.PointOfInterest.id.in_(existing_poi_ids)) ) pois = (await db.scalars(poi_query)).all() poi_map = {poi.id: poi for poi in pois} order = 0 for item in track: poi_object = None if item.id: poi_object = poi_map.get(item.id) if not poi_object: poi_object = await models.PointOfInterest.create(db, data=dict(created_by_id=user_id, name=item.name)) await db.flush() await models.FlightTrack.create( db, data={ "flight_id": flight.id, "point_of_interest_id": poi_object.id, "order": order } ) order += 1 async def handle_copilot_edit(db: AsyncSession, copilot: CopilotInput, user_id: int) -> int: if copilot.id: return copilot.id else: copilot = await models.Copilot.create( db, data={ "name": copilot.name, "created_by_id": user_id, } ) await db.flush() return copilot.id @strawberry.type class EditFlightMutation: @strawberry_sqlalchemy_input(models.Flight, exclude_fields=["id", "copilot_id", "deleted"], all_optional=True) class EditFlightInput: track: Optional[List[PointOfInterestInput]] = None copilot: Optional[CopilotInput] = None @strawberry.mutation async def edit_flight(self, info, id: int, input: EditFlightInput) -> Flight: # TODO: umoznit editovat jen vlastni lety! flight = await models.Flight.update(info.context.db, id=id, data=input.to_dict()) if input.track is not None: await handle_track_edit(db=info.context.db, flight=flight, track=input.track, user_id=info.context.user_id) if flight.solo: flight.copilot_id = None elif input.copilot is not None: flight.copilot_id = await handle_copilot_edit(info.context.db, input.copilot, info.context.user_id) return flight @strawberry.type class DeleteFlightMutation: @strawberry.mutation async def delete_flight(self, info, id: int) -> Flight: flight = ( (await info.context.db.scalars( get_base_query(info.context.user_id) .filter(models.Flight.id == id)) ) .one() ) flight.deleted = True return flight