47 lines
1.5 KiB
Python
47 lines
1.5 KiB
Python
from __future__ import annotations
|
|
|
|
from uuid import UUID
|
|
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.orm import Session
|
|
|
|
from backend.app.db.models import TrainSchedule
|
|
from backend.app.models import TrainScheduleCreate
|
|
from backend.app.repositories.base import BaseRepository
|
|
|
|
|
|
class TrainScheduleRepository(BaseRepository[TrainSchedule]):
|
|
"""Persistence operations for train timetables."""
|
|
|
|
model = TrainSchedule
|
|
|
|
def __init__(self, session: Session) -> None:
|
|
super().__init__(session)
|
|
|
|
@staticmethod
|
|
def _ensure_uuid(value: UUID | str) -> UUID:
|
|
if isinstance(value, UUID):
|
|
return value
|
|
return UUID(str(value))
|
|
|
|
def list_for_train(self, train_id: UUID | str) -> list[TrainSchedule]:
|
|
identifier = self._ensure_uuid(train_id)
|
|
statement = (
|
|
sa.select(self.model)
|
|
.where(self.model.train_id == identifier)
|
|
.order_by(self.model.sequence_index.asc())
|
|
)
|
|
return list(self.session.scalars(statement))
|
|
|
|
def create(self, data: TrainScheduleCreate) -> TrainSchedule:
|
|
schedule = TrainSchedule(
|
|
train_id=self._ensure_uuid(data.train_id),
|
|
station_id=self._ensure_uuid(data.station_id),
|
|
sequence_index=data.sequence_index,
|
|
scheduled_arrival=data.scheduled_arrival,
|
|
scheduled_departure=data.scheduled_departure,
|
|
dwell_seconds=data.dwell_seconds,
|
|
)
|
|
self.session.add(schedule)
|
|
return schedule
|