Model Mapping with GeoAlchemy2

-- Run in the target database; must return a version string
SELECT PostGIS_Version();

-- If the extension is missing, enable it:
CREATE EXTENSION IF NOT EXISTS postgis;

pip install "sqlalchemy>=2.0,<3" "geoalchemy2>=0.14" "psycopg2-binary>=2.9" alembic

-- Create a temporary table and confirm the GiST index type is available
CREATE TEMP TABLE _gist_check (geom geometry(POINT, 4326));
CREATE INDEX ON _gist_check USING gist (geom);
DROP TABLE _gist_check;
-- No error means the PostGIS operator classes are installed correctly

from sqlalchemy import create_engine, URL
from sqlalchemy.orm import DeclarativeBase, sessionmaker

engine = create_engine(
    URL.create(
        drivername="postgresql+psycopg2",
        username="app_user",
        password="secure_password",
        host="localhost",
        port=5432,
        database="spatial_db",
    ),
    pool_size=10,
    max_overflow=20,
    pool_pre_ping=True,   # validates connections; avoids stale-socket errors
    echo=False,           # set True during development to inspect compiled SQL
)

class Base(DeclarativeBase):
    pass

SessionLocal = sessionmaker(
    bind=engine,
    expire_on_commit=False,  # geometry columns remain accessible after commit
)

from sqlalchemy import String, Integer, Index
from sqlalchemy.orm import Mapped, mapped_column
from geoalchemy2 import Geometry

class Facility(Base):
    __tablename__ = "facilities"
    __table_args__ = (
        # GiST index is required for ST_DWithin, ST_Intersects, and
        # bounding-box operators to use index scans rather than seq scans
        Index("idx_facilities_location", "location", postgresql_using="gist"),
        {"schema": "public"},
    )

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    name: Mapped[str] = mapped_column(String(255), nullable=False)
    category: Mapped[str] = mapped_column(String(50), index=True)

    # Explicit SRID=4326 (WGS84). PostGIS stores this in geometry_columns
    # and enforces it via a CHECK constraint on insert.
    location: Mapped[Geometry] = mapped_column(
        Geometry(geometry_type="POINT", srid=4326),
        nullable=False,
    )

class Zone(Base):
    __tablename__ = "zones"
    __table_args__ = (
        Index("idx_zones_boundary", "boundary", postgresql_using="gist"),
        Index("idx_zones_centroid", "centroid", postgresql_using="gist"),
    )

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    boundary: Mapped[Geometry] = mapped_column(
        Geometry(geometry_type="POLYGON", srid=4326), nullable=False
    )
    centroid: Mapped[Geometry] = mapped_column(
        Geometry(geometry_type="POINT", srid=4326), nullable=True
    )

from geoalchemy2.elements import WKTElement
from sqlalchemy.orm import Session

def create_facility(
    name: str, category: str, lat: float, lon: float
) -> Facility:
    # WKTElement carries SRID alongside the coordinate string
    geom = WKTElement(f"POINT({lon} {lat})", srid=4326)

    with Session(engine) as session:
        facility = Facility(name=name, category=category, location=geom)
        session.add(facility)
        session.commit()
        session.refresh(facility)  # loads server-generated id and defaults
        return facility

from shapely.geometry import Point

facility.location = Point(-73.9857, 40.7484)
# GeoAlchemy2 converts to WKBElement with the column's srid=4326 on flush

from sqlalchemy import select, func
from geoalchemy2.elements import WKTElement

def facilities_within_km(lat: float, lon: float, km: float) -> list[Facility]:
    origin = WKTElement(f"POINT({lon} {lat})", srid=4326)

    stmt = (
        select(Facility)
        .where(
            # ST_DWithin on geography uses metres; cast geometry → geography
            func.ST_DWithin(
                Facility.location.cast("geography"),
                func.ST_GeogFromText(f"POINT({lon} {lat})"),
                km * 1000,  # metres
            )
        )
        .order_by(
            func.ST_Distance(
                Facility.location.cast("geography"),
                func.ST_GeogFromText(f"POINT({lon} {lat})"),
            )
        )
        .limit(50)
    )

    with Session(engine) as session:
        return session.scalars(stmt).all()

from geoalchemy2.shape import to_shape
from shapely.geometry import mapping

def serialize_facility(facility: Facility) -> dict:
    geom = to_shape(facility.location)  # WKBElement → Shapely geometry
    return {
        "id": facility.id,
        "name": facility.name,
        "category": facility.category,
        "coordinates": list(geom.coords),  # [(lon, lat)]
        "geojson": mapping(geom),           # {"type": "Point", "coordinates": [...]}
    }

from sqlalchemy.ext.hybrid import hybrid_property
from sqlalchemy import func
from geoalchemy2.shape import to_shape

class Facility(Base):
    # ... columns as above ...

    @hybrid_property
    def centroid(self):
        """Python-side: returns a Shapely Point."""
        if self.location is None:
            return None
        return to_shape(self.location).centroid

    @centroid.expression
    def centroid(cls):
        """SQL-side: compiles to ST_Centroid(location)."""
        return func.ST_Centroid(cls.location)

    @hybrid_property
    def longitude(self):
        if self.location is None:
            return None
        return to_shape(self.location).x

    @longitude.expression
    def longitude(cls):
        return func.ST_X(cls.location)

from sqlalchemy import event, text

@event.listens_for(engine, "connect")
def set_spatial_gucs(dbapi_conn, _):
    with dbapi_conn.cursor() as cur:
        cur.execute("SET random_page_cost = 1.1")
        cur.execute("SET work_mem = '32MB'")

-- Check what SRID the column expects
SELECT type, srid
FROM geometry_columns
WHERE f_table_name = 'facilities' AND f_geometry_column = 'location';

# Wrong — SRID defaults to 0
location = WKTElement("POINT(-73.98 40.74)")

# Correct
location = WKTElement("POINT(-73.98 40.74)", srid=4326)

SELECT indexname, indexdef
FROM pg_indexes
WHERE tablename = 'facilities'
  AND indexdef ILIKE '%gist%';
-- Empty result means no GiST index exists

CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_facilities_location
ON facilities USING gist (location);

-- Then refresh statistics
ANALYZE facilities;

SessionLocal = sessionmaker(bind=engine, expire_on_commit=False)

session.commit()
session.refresh(facility)   # reloads id, location, and any server defaults
print(facility.location)    # now populated

stmt = select(Facility).execution_options(yield_per=500)
with Session(engine) as session:
    for batch in session.scalars(stmt).partitions():
        process(batch)

-- 1. Confirm the geometry column is registered with correct type and SRID
SELECT f_table_name, f_geometry_column, type, srid
FROM geometry_columns
WHERE f_table_name = 'facilities';
-- Expected: facilities | location | POINT | 4326

-- 2. Confirm the GiST index exists
SELECT indexname
FROM pg_indexes
WHERE tablename = 'facilities' AND indexdef ILIKE '%gist%';
-- Expected: idx_facilities_location

-- 3. Insert a test row and confirm it round-trips correctly
INSERT INTO facilities (name, category, location)
VALUES ('Test', 'qa', ST_GeomFromText('POINT(-73.98 40.74)', 4326));

SELECT id, name, ST_AsText(location) AS wkt, ST_SRID(location) AS srid
FROM facilities
WHERE name = 'Test';
-- Expected: wkt=POINT(-73.98 40.74), srid=4326

-- 4. Check the query uses the index
EXPLAIN (ANALYZE, BUFFERS)
SELECT * FROM facilities
WHERE ST_DWithin(location::geography, ST_GeogFromText('POINT(-73.98 40.74)'), 5000);
-- Expected: "Index Scan using idx_facilities_location"

-- Cleanup
DELETE FROM facilities WHERE name = 'Test';