145 lines
4.7 KiB
Python
145 lines
4.7 KiB
Python
from typing import Iterable, Self
|
|
|
|
from sqlalchemy import (
|
|
BigInteger,
|
|
Column,
|
|
Enum,
|
|
Float,
|
|
ForeignKey,
|
|
select,
|
|
String,
|
|
Table,
|
|
)
|
|
from sqlalchemy.orm import Mapped, relationship, selectinload, with_polymorphic
|
|
from sqlalchemy.schema import Index
|
|
|
|
from ..db import Base, db
|
|
from ..idfm_interface.idfm_types import TransportMode, IdfmState, StopAreaType
|
|
|
|
stop_area_stop_association_table = Table(
|
|
"stop_area_stop_association_table",
|
|
Base.metadata,
|
|
Column("stop_id", ForeignKey("_stops.id")),
|
|
Column("stop_area_id", ForeignKey("stop_areas.id")),
|
|
)
|
|
|
|
|
|
class _Stop(Base):
|
|
|
|
db = db
|
|
|
|
id = Column(BigInteger, primary_key=True)
|
|
kind = Column(String)
|
|
|
|
name = Column(String, nullable=False, index=True)
|
|
town_name = Column(String, nullable=False)
|
|
postal_region = Column(String, nullable=False)
|
|
xepsg2154 = Column(BigInteger, nullable=False)
|
|
yepsg2154 = Column(BigInteger, nullable=False)
|
|
version = Column(String, nullable=False)
|
|
created_ts = Column(BigInteger)
|
|
changed_ts = Column(BigInteger, nullable=False)
|
|
lines: Mapped[list["Line"]] = relationship(
|
|
"Line",
|
|
secondary="line_stop_association_table",
|
|
back_populates="stops",
|
|
# lazy="joined",
|
|
lazy="selectin",
|
|
)
|
|
areas: Mapped[list["StopArea"]] = relationship(
|
|
"StopArea", secondary=stop_area_stop_association_table, back_populates="stops"
|
|
)
|
|
|
|
__tablename__ = "_stops"
|
|
__mapper_args__ = {"polymorphic_identity": "_stops", "polymorphic_on": kind}
|
|
__table_args__ = (
|
|
# To optimize the ilike requests
|
|
Index(
|
|
"name_idx_gin",
|
|
name,
|
|
postgresql_ops={"name": "gin_trgm_ops"},
|
|
postgresql_using="gin",
|
|
),
|
|
)
|
|
|
|
# TODO: Test https://www.cybertec-postgresql.com/en/postgresql-more-performance-for-like-and-ilike-statements/
|
|
# TODO: Should be able to remove with_polymorphic ?
|
|
@classmethod
|
|
async def get_by_name(cls, name: str) -> list[Self]:
|
|
stop_stop_area = with_polymorphic(_Stop, [Stop, StopArea])
|
|
stmt = (
|
|
select(stop_stop_area)
|
|
.where(stop_stop_area.name.ilike(f"%{name}%"))
|
|
.options(
|
|
selectinload(stop_stop_area.areas),
|
|
selectinload(stop_stop_area.lines),
|
|
)
|
|
)
|
|
res = await cls.db.session.execute(stmt)
|
|
return res.scalars()
|
|
|
|
|
|
class Stop(_Stop):
|
|
|
|
id = Column(BigInteger, ForeignKey("_stops.id"), primary_key=True)
|
|
|
|
latitude = Column(Float, nullable=False)
|
|
longitude = Column(Float, nullable=False)
|
|
transport_mode = Column(Enum(TransportMode), nullable=False)
|
|
accessibility = Column(Enum(IdfmState), nullable=False)
|
|
visual_signs_available = Column(Enum(IdfmState), nullable=False)
|
|
audible_signs_available = Column(Enum(IdfmState), nullable=False)
|
|
record_id = Column(String, nullable=False)
|
|
record_ts = Column(BigInteger, nullable=False)
|
|
|
|
__tablename__ = "stops"
|
|
__mapper_args__ = {"polymorphic_identity": "stops", "polymorphic_load": "inline"}
|
|
|
|
|
|
class StopArea(_Stop):
|
|
|
|
id = Column(BigInteger, ForeignKey("_stops.id"), primary_key=True)
|
|
|
|
type = Column(Enum(StopAreaType), nullable=False)
|
|
stops: Mapped[list[_Stop]] = relationship(
|
|
_Stop,
|
|
secondary=stop_area_stop_association_table,
|
|
back_populates="areas",
|
|
lazy="selectin",
|
|
# lazy="joined",
|
|
)
|
|
|
|
__tablename__ = "stop_areas"
|
|
__mapper_args__ = {"polymorphic_identity": "stop_areas", "polymorphic_load": "inline"}
|
|
|
|
@classmethod
|
|
async def add_stops(cls, stop_area_to_stop_ids: Iterable[tuple[str, str]]) -> int:
|
|
stop_area_ids, stop_ids = set(), set()
|
|
for stop_area_id, stop_id in stop_area_to_stop_ids:
|
|
stop_area_ids.add(stop_area_id)
|
|
stop_ids.add(stop_id)
|
|
|
|
res = await cls.db.session.execute(
|
|
select(StopArea)
|
|
.where(StopArea.id.in_(stop_area_ids))
|
|
.options(selectinload(StopArea.stops))
|
|
)
|
|
stop_areas = {stop_area.id: stop_area for stop_area in res.scalars()}
|
|
|
|
res = await cls.db.session.execute(select(_Stop).where(_Stop.id.in_(stop_ids)))
|
|
stops = {stop.id: stop for stop in res.scalars()}
|
|
|
|
found = 0
|
|
for stop_area_id, stop_id in stop_area_to_stop_ids:
|
|
if (stop_area := stop_areas.get(stop_area_id)) is not None:
|
|
if (stop := stops.get(stop_id)) is not None:
|
|
stop_area.stops.append(stop)
|
|
found += 1
|
|
else:
|
|
print(f"No stop found for {stop_id} id")
|
|
else:
|
|
print(f"No stop area found for {stop_area_id}")
|
|
|
|
await cls.db.session.commit()
|
|
return found
|