🎉 First commit !!!

backend/idfm_matrix_backend/models/stop.py

@@ -0,0 +1,144 @@
+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