42 lines
1.2 KiB
Python
42 lines
1.2 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 Train
|
|
from backend.app.models import TrainCreate
|
|
from backend.app.repositories.base import BaseRepository
|
|
|
|
|
|
class TrainRepository(BaseRepository[Train]):
|
|
model = Train
|
|
|
|
def __init__(self, session: Session) -> None:
|
|
super().__init__(session)
|
|
|
|
def list_all(self) -> list[Train]:
|
|
statement = sa.select(self.model)
|
|
return list(self.session.scalars(statement))
|
|
|
|
@staticmethod
|
|
def _optional_uuid(value: UUID | str | None) -> UUID | None:
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, UUID):
|
|
return value
|
|
return UUID(str(value))
|
|
|
|
def create(self, data: TrainCreate) -> Train:
|
|
train = Train(
|
|
designation=data.designation,
|
|
operator_id=self._optional_uuid(data.operator_id),
|
|
home_station_id=self._optional_uuid(data.home_station_id),
|
|
capacity=data.capacity,
|
|
max_speed_kph=data.max_speed_kph,
|
|
consist=data.consist,
|
|
)
|
|
self.session.add(train)
|
|
return train
|