🎨 Reorganize back-end code
This commit is contained in:
275
backend/api/models/stop.py
Normal file
275
backend/api/models/stop.py
Normal file
@@ -0,0 +1,275 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from logging import getLogger
|
||||
from typing import Iterable, Sequence, TYPE_CHECKING
|
||||
|
||||
from sqlalchemy import (
|
||||
BigInteger,
|
||||
Computed,
|
||||
desc,
|
||||
Enum,
|
||||
Float,
|
||||
ForeignKey,
|
||||
func,
|
||||
Integer,
|
||||
JSON,
|
||||
select,
|
||||
String,
|
||||
)
|
||||
from sqlalchemy.orm import (
|
||||
mapped_column,
|
||||
Mapped,
|
||||
relationship,
|
||||
selectinload,
|
||||
with_polymorphic,
|
||||
)
|
||||
from sqlalchemy.schema import Index
|
||||
from sqlalchemy_utils.types.ts_vector import TSVectorType
|
||||
|
||||
from db import Base, db
|
||||
from idfm_interface.idfm_types import TransportMode, IdfmState, StopAreaType
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .line import Line
|
||||
|
||||
|
||||
logger = getLogger(__name__)
|
||||
|
||||
class StopAreaStopAssociations(Base):
|
||||
|
||||
id = mapped_column(BigInteger, primary_key=True)
|
||||
stop_id = mapped_column(BigInteger, ForeignKey("_stops.id"))
|
||||
stop_area_id = mapped_column(BigInteger, ForeignKey("stop_areas.id"))
|
||||
|
||||
__tablename__ = "stop_area_stop_associations"
|
||||
|
||||
|
||||
class _Stop(Base):
|
||||
|
||||
db = db
|
||||
|
||||
id = mapped_column(BigInteger, primary_key=True)
|
||||
kind = mapped_column(String)
|
||||
|
||||
name = mapped_column(String, nullable=False, index=True)
|
||||
town_name = mapped_column(String, nullable=False)
|
||||
postal_region = mapped_column(Integer, nullable=False)
|
||||
epsg3857_x = mapped_column(Float, nullable=False)
|
||||
epsg3857_y = mapped_column(Float, nullable=False)
|
||||
|
||||
version = mapped_column(String, nullable=False)
|
||||
created_ts = mapped_column(BigInteger)
|
||||
changed_ts = mapped_column(BigInteger, nullable=False)
|
||||
|
||||
lines: Mapped[list[Line]] = relationship(
|
||||
"Line",
|
||||
secondary="line_stop_associations",
|
||||
back_populates="stops",
|
||||
lazy="selectin",
|
||||
)
|
||||
areas: Mapped[list["StopArea"]] = relationship(
|
||||
"StopArea",
|
||||
secondary="stop_area_stop_associations",
|
||||
back_populates="stops",
|
||||
)
|
||||
connection_area_id: Mapped[int] = mapped_column(
|
||||
ForeignKey("connection_areas.id"), nullable=True
|
||||
)
|
||||
connection_area: Mapped["ConnectionArea"] = relationship(
|
||||
back_populates="stops", lazy="selectin"
|
||||
)
|
||||
|
||||
names_tsv = mapped_column(
|
||||
TSVectorType("name", "town_name", regconfig="french"),
|
||||
Computed("to_tsvector('french', name || ' ' || town_name)", persisted=True),
|
||||
)
|
||||
|
||||
__tablename__ = "_stops"
|
||||
__mapper_args__ = {"polymorphic_identity": "_stops", "polymorphic_on": kind}
|
||||
__table_args__ = (
|
||||
Index(
|
||||
"names_tsv_idx",
|
||||
names_tsv,
|
||||
postgresql_ops={"name": "gin_trgm_ops"},
|
||||
postgresql_using="gin",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def get_by_name(cls, name: str) -> Sequence[_Stop] | None:
|
||||
if (session := await cls.db.get_session()) is not None:
|
||||
|
||||
async with session.begin():
|
||||
descendants = with_polymorphic(_Stop, "*")
|
||||
|
||||
match_stmt = descendants.names_tsv.match(
|
||||
name, postgresql_regconfig="french"
|
||||
)
|
||||
ranking_stmt = func.ts_rank_cd(
|
||||
descendants.names_tsv, func.plainto_tsquery("french", name)
|
||||
)
|
||||
stmt = (
|
||||
select(descendants).filter(match_stmt).order_by(desc(ranking_stmt))
|
||||
)
|
||||
|
||||
res = await session.execute(stmt)
|
||||
stops = res.scalars().all()
|
||||
|
||||
return stops
|
||||
|
||||
return None
|
||||
|
||||
|
||||
class Stop(_Stop):
|
||||
|
||||
id = mapped_column(BigInteger, ForeignKey("_stops.id"), primary_key=True)
|
||||
|
||||
transport_mode = mapped_column(Enum(TransportMode), nullable=False)
|
||||
accessibility = mapped_column(Enum(IdfmState), nullable=False)
|
||||
visual_signs_available = mapped_column(Enum(IdfmState), nullable=False)
|
||||
audible_signs_available = mapped_column(Enum(IdfmState), nullable=False)
|
||||
|
||||
record_id = mapped_column(String, nullable=False)
|
||||
record_ts = mapped_column(BigInteger, nullable=False)
|
||||
|
||||
__tablename__ = "stops"
|
||||
__mapper_args__ = {"polymorphic_identity": "stops", "polymorphic_load": "inline"}
|
||||
|
||||
|
||||
class StopArea(_Stop):
|
||||
|
||||
id = mapped_column(BigInteger, ForeignKey("_stops.id"), primary_key=True)
|
||||
|
||||
type = mapped_column(Enum(StopAreaType), nullable=False)
|
||||
|
||||
stops: Mapped[list["Stop"]] = relationship(
|
||||
"Stop",
|
||||
secondary="stop_area_stop_associations",
|
||||
back_populates="areas",
|
||||
lazy="selectin",
|
||||
)
|
||||
|
||||
__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[int, int]]
|
||||
) -> int | None:
|
||||
if (session := await cls.db.get_session()) is not None:
|
||||
|
||||
async with session.begin():
|
||||
|
||||
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)
|
||||
|
||||
stop_areas_res = await session.scalars(
|
||||
select(StopArea)
|
||||
.where(StopArea.id.in_(stop_area_ids))
|
||||
.options(selectinload(StopArea.stops))
|
||||
)
|
||||
stop_areas: dict[int, StopArea] = {
|
||||
stop_area.id: stop_area for stop_area in stop_areas_res.all()
|
||||
}
|
||||
|
||||
stop_res = await session.execute(
|
||||
select(Stop).where(Stop.id.in_(stop_ids))
|
||||
)
|
||||
stops: dict[int, Stop] = {stop.id: stop for stop in stop_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}")
|
||||
|
||||
return found
|
||||
|
||||
return None
|
||||
|
||||
|
||||
class StopShape(Base):
|
||||
|
||||
db = db
|
||||
|
||||
id = mapped_column(BigInteger, primary_key=True) # Same id than ConnectionArea
|
||||
type = mapped_column(Integer, nullable=False)
|
||||
epsg3857_bbox = mapped_column(JSON)
|
||||
epsg3857_points = mapped_column(JSON)
|
||||
|
||||
__tablename__ = "stop_shapes"
|
||||
|
||||
|
||||
class ConnectionArea(Base):
|
||||
|
||||
db = db
|
||||
|
||||
id = mapped_column(BigInteger, primary_key=True)
|
||||
|
||||
name = mapped_column(String, nullable=False)
|
||||
town_name = mapped_column(String, nullable=False)
|
||||
postal_region = mapped_column(String, nullable=False)
|
||||
epsg3857_x = mapped_column(Float, nullable=False)
|
||||
epsg3857_y = mapped_column(Float, nullable=False)
|
||||
transport_mode = mapped_column(Enum(StopAreaType), nullable=False)
|
||||
|
||||
version = mapped_column(String, nullable=False)
|
||||
created_ts = mapped_column(BigInteger)
|
||||
changed_ts = mapped_column(BigInteger, nullable=False)
|
||||
|
||||
stops: Mapped[list["_Stop"]] = relationship(back_populates="connection_area")
|
||||
|
||||
__tablename__ = "connection_areas"
|
||||
|
||||
# TODO: Merge with StopArea.add_stops
|
||||
@classmethod
|
||||
async def add_stops(
|
||||
cls, conn_area_to_stop_ids: Iterable[tuple[int, int]]
|
||||
) -> int | None:
|
||||
if (session := await cls.db.get_session()) is not None:
|
||||
|
||||
async with session.begin():
|
||||
|
||||
conn_area_ids, stop_ids = set(), set()
|
||||
for conn_area_id, stop_id in conn_area_to_stop_ids:
|
||||
conn_area_ids.add(conn_area_id)
|
||||
stop_ids.add(stop_id)
|
||||
|
||||
conn_area_res = await session.execute(
|
||||
select(ConnectionArea)
|
||||
.where(ConnectionArea.id.in_(conn_area_ids))
|
||||
.options(selectinload(ConnectionArea.stops))
|
||||
)
|
||||
conn_areas: dict[int, ConnectionArea] = {
|
||||
conn.id: conn for conn in conn_area_res.scalars()
|
||||
}
|
||||
|
||||
stop_res = await session.execute(
|
||||
select(Stop).where(Stop.id.in_(stop_ids))
|
||||
)
|
||||
stops: dict[int, Stop] = {stop.id: stop for stop in stop_res.scalars()}
|
||||
|
||||
found = 0
|
||||
for conn_area_id, stop_id in conn_area_to_stop_ids:
|
||||
if (conn_area := conn_areas.get(conn_area_id)) is not None:
|
||||
if (stop := stops.get(stop_id)) is not None:
|
||||
conn_area.stops.append(stop)
|
||||
found += 1
|
||||
else:
|
||||
print(f"No stop found for {stop_id} id")
|
||||
else:
|
||||
print(f"No connection area found for {conn_area_id}")
|
||||
|
||||
return found
|
||||
|
||||
return None
|
||||
Reference in New Issue
Block a user