52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
from collections import defaultdict
|
|
from typing import List, Optional
|
|
from sqlalchemy import select
|
|
from strawberry.dataloader import DataLoader
|
|
from database import async_session
|
|
from database.models import Flight, Copilot, PointOfInterest, Event
|
|
|
|
|
|
class FlightsLoader:
|
|
def __init__(self, relationship_column, extra_join: Optional[list] = None):
|
|
if extra_join is None:
|
|
extra_join = []
|
|
|
|
self.relationship_column = relationship_column
|
|
self.extra_join = extra_join
|
|
|
|
async def load(self, ids: List[int]):
|
|
async with async_session() as session:
|
|
rel_column = self.relationship_column
|
|
query = (
|
|
select(Flight, rel_column)
|
|
.filter(rel_column.in_(ids))
|
|
.order_by(Flight.takeoff_datetime.desc())
|
|
)
|
|
|
|
for table in self.extra_join:
|
|
query = query.join(table)
|
|
|
|
data = (await session.execute(query)).all()
|
|
|
|
result_data = defaultdict(list)
|
|
for item, rel_id in data:
|
|
result_data[rel_id].append(item)
|
|
|
|
return [result_data[id_] for id_ in ids]
|
|
|
|
|
|
flights_by_copilot_dataloader = DataLoader(
|
|
load_fn=FlightsLoader(Copilot.id, extra_join=[Flight.copilots]).load,
|
|
cache=False
|
|
)
|
|
flights_by_aircraft_dataloader = DataLoader(load_fn=FlightsLoader(Flight.aircraft_id).load, cache=False)
|
|
flight_by_poi_dataloader = DataLoader(
|
|
load_fn=FlightsLoader(PointOfInterest.id, extra_join=[Flight.track, PointOfInterest]).load,
|
|
cache=False
|
|
)
|
|
|
|
flights_by_event_dataloader = DataLoader(
|
|
load_fn=FlightsLoader(Event.id, extra_join=[Flight.event]).load,
|
|
cache=False
|
|
)
|