models.py

"""SQLAlchemy 2.0 models for a PostGIS schema.

Generated by geomodelgen from schema 'public'. Edits will be lost the next
time the generator runs; change the database (or the schema description) and
regenerate instead.
"""

from __future__ import annotations

import datetime
import decimal
import enum
import typing
import uuid
from typing import Optional

from geoalchemy2 import Geography, Geometry
from geoalchemy2.elements import WKBElement
from sqlalchemy import (
    ARRAY,
    BigInteger,
    Boolean,
    CheckConstraint,
    Date,
    DateTime,
    Enum,
    Float,
    ForeignKey,
    Index,
    Integer,
    Numeric,
    SmallInteger,
    String,
    Text,
    UniqueConstraint,
    Uuid,
    text,
)
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship


class Base(DeclarativeBase):
    """Declarative base shared by every generated model."""


class RoadSurface(enum.Enum):
    """PostgreSQL enum type 'road_surface'."""

    asphalt = "asphalt"
    concrete = "concrete"
    gravel = "gravel"
    unpaved = "unpaved"


class ServiceTier(enum.Enum):
    """PostgreSQL enum type 'service_tier'."""

    standard = "standard"
    priority = "priority"
    emergency = "emergency"


class City(Base):
    """Populated places, one point per city centre."""

    __tablename__ = "cities"

    __table_args__ = (
        Index("cities_geom_idx", "geom", postgresql_using="gist"),
        Index("cities_name_idx", "name", postgresql_using="btree"),
        UniqueConstraint("region_id", "name", name="cities_region_id_name_key"),
        CheckConstraint("population >= 0", name="cities_population_check"),
        {"comment": "Populated places, one point per city centre."},
    )

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
    region_id: Mapped[Optional[int]] = mapped_column(
        Integer, ForeignKey("regions.id", ondelete="SET NULL")
    )
    name: Mapped[str] = mapped_column(Text, nullable=False)
    population: Mapped[Optional[int]] = mapped_column(Integer)
    area_km2: Mapped[Optional[decimal.Decimal]] = mapped_column(Numeric(12, 3))
    tags: Mapped[Optional[list[str]]] = mapped_column(ARRAY(Text))
    properties: Mapped[Optional[dict[str, typing.Any]]] = mapped_column(
        JSONB, server_default=text("'{}'::jsonb")
    )
    external_id: Mapped[Optional[uuid.UUID]] = mapped_column(Uuid)
    geom: Mapped[WKBElement] = mapped_column(
        Geometry(geometry_type="POINT", srid=4326, spatial_index=False),
        nullable=False,
        comment="City centre.",
    )
    updated_at: Mapped[datetime.datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=text("now()")
    )

    region: Mapped[Optional["Region"]] = relationship("Region", back_populates="cities")
    roads: Mapped[list["Road"]] = relationship("Road", back_populates="city")


class Region(Base):
    """Administrative regions used to group cities and service areas."""

    __tablename__ = "regions"

    __table_args__ = (
        UniqueConstraint("code", name="regions_code_key"),
        {"comment": "Administrative regions used to group cities and service areas."},
    )

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    code: Mapped[str] = mapped_column(
        String(8), nullable=False, comment="ISO 3166-2 subdivision code."
    )
    name: Mapped[str] = mapped_column(Text, nullable=False)
    created_at: Mapped[datetime.datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=text("now()")
    )

    cities: Mapped[list["City"]] = relationship("City", back_populates="region")
    service_areas: Mapped[list["ServiceArea"]] = relationship(
        "ServiceArea", back_populates="region"
    )


class Road(Base):
    """Road centrelines with a 3D geometry."""

    __tablename__ = "roads"

    __table_args__ = (
        Index("roads_geom_idx", "geom", postgresql_using="gist"),
        Index(
            "roads_toll_geom_idx", "geom", postgresql_using="gist", postgresql_where=text("is_toll")
        ),
        CheckConstraint("lanes > 0 AND lanes < 12", name="roads_lanes_check"),
        {"comment": "Road centrelines with a 3D geometry."},
    )

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
    city_id: Mapped[int] = mapped_column(
        BigInteger, ForeignKey("cities.id", ondelete="CASCADE"), nullable=False
    )
    name: Mapped[Optional[str]] = mapped_column(Text)
    surface: Mapped[RoadSurface] = mapped_column(
        Enum(RoadSurface, name="road_surface", native_enum=True),
        nullable=False,
        server_default=text("'asphalt'::road_surface"),
    )
    lanes: Mapped[Optional[int]] = mapped_column(SmallInteger)
    speed_limit_kph: Mapped[Optional[int]] = mapped_column(SmallInteger)
    is_toll: Mapped[bool] = mapped_column(Boolean, nullable=False, server_default=text("false"))
    geom: Mapped[WKBElement] = mapped_column(
        Geometry(geometry_type="LINESTRINGZ", srid=4326, dimension=3, spatial_index=False),
        nullable=False,
    )

    city: Mapped["City"] = relationship("City", back_populates="roads")


class ServiceArea(Base):
    """Polygonal coverage areas, stored as geography for metre-accurate distance."""

    __tablename__ = "service_areas"

    __table_args__ = (
        {"comment": "Polygonal coverage areas, stored as geography for metre-accurate distance."},
    )

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    region_id: Mapped[int] = mapped_column(
        Integer, ForeignKey("regions.id", ondelete="CASCADE"), nullable=False
    )
    label: Mapped[str] = mapped_column(Text, nullable=False)
    tier: Mapped[ServiceTier] = mapped_column(
        Enum(ServiceTier, name="service_tier", native_enum=True),
        nullable=False,
        server_default=text("'standard'::service_tier"),
    )
    valid_from: Mapped[Optional[datetime.date]] = mapped_column(Date)
    boundary: Mapped[WKBElement] = mapped_column(
        Geography(geometry_type="POLYGON", srid=4326, spatial_index=True), nullable=False
    )

    region: Mapped["Region"] = relationship("Region", back_populates="service_areas")


class SurveyPoint(Base):
    """Raw survey captures; deliberately unindexed to show the missing-index warning."""

    __tablename__ = "survey_points"

    __table_args__ = (
        {
            "comment": "Raw survey captures; deliberately unindexed to show the missing-index warning."
        },
    )

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
    captured_at: Mapped[datetime.datetime] = mapped_column(DateTime(timezone=True), nullable=False)
    accuracy_m: Mapped[Optional[float]] = mapped_column(Float)
    location: Mapped[WKBElement] = mapped_column(
        Geometry(geometry_type="POINTZM", srid=4326, dimension=4, spatial_index=False),
        nullable=False,
    )  # no spatial index on this column in the source schema; see queries.py for the CREATE INDEX to run