Loading GTFS into PostGIS with Python

import io
import logging
from pathlib import Path

import pandas as pd
import sqlalchemy as sa
from sqlalchemy import text

logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")

ENGINE = sa.create_engine(
    "postgresql+psycopg2://gtfs:gtfs@localhost:5432/transit", future=True
)

DDL = """
CREATE EXTENSION IF NOT EXISTS postgis;

CREATE TABLE IF NOT EXISTS stops (
    stop_id   TEXT PRIMARY KEY,
    stop_name TEXT,
    stop_lat  DOUBLE PRECISION NOT NULL,
    stop_lon  DOUBLE PRECISION NOT NULL
);

CREATE TABLE IF NOT EXISTS trips (
    trip_id    TEXT PRIMARY KEY,
    route_id   TEXT NOT NULL,
    service_id TEXT NOT NULL,
    shape_id   TEXT
);

CREATE TABLE IF NOT EXISTS stop_times (
    trip_id        TEXT NOT NULL,
    stop_id        TEXT NOT NULL,
    stop_sequence  INTEGER NOT NULL,
    arrival_time   TEXT,
    departure_time TEXT,
    PRIMARY KEY (trip_id, stop_sequence)
);

CREATE TABLE IF NOT EXISTS shapes (
    shape_id          TEXT NOT NULL,
    shape_pt_lat      DOUBLE PRECISION NOT NULL,
    shape_pt_lon      DOUBLE PRECISION NOT NULL,
    shape_pt_sequence INTEGER NOT NULL,
    PRIMARY KEY (shape_id, shape_pt_sequence)
);

CREATE TABLE IF NOT EXISTS shape_geom (
    shape_id TEXT PRIMARY KEY,
    geom     geometry(LineString, 4326)
);
"""

# Explicit dtypes: IDs stay strings, sequences are ints, times are text.
DTYPES = {
    "stops": {"stop_id": str, "stop_name": str,
              "stop_lat": float, "stop_lon": float},
    "trips": {"trip_id": str, "route_id": str,
              "service_id": str, "shape_id": str},
    "stop_times": {"trip_id": str, "stop_id": str, "stop_sequence": "Int32",
                   "arrival_time": str, "departure_time": str},
    "shapes": {"shape_id": str, "shape_pt_lat": float, "shape_pt_lon": float,
               "shape_pt_sequence": "Int32"},
}


def create_schema() -> None:
    with ENGINE.begin() as conn:
        for statement in filter(str.strip, DDL.split(";")):
            conn.execute(text(statement))
    logging.info("Schema created")


def copy_table(feed_dir: Path, table: str) -> int:
    """Load one GTFS file into its table via COPY, with explicit dtypes."""
    dtypes = DTYPES[table]
    frame = pd.read_csv(
        feed_dir / f"{table}.txt", dtype=dtypes, usecols=list(dtypes)
    )
    buffer = io.StringIO()
    frame.to_csv(buffer, index=False, header=False, na_rep="")
    buffer.seek(0)
    columns = ", ".join(frame.columns)
    raw = ENGINE.raw_connection()
    try:
        with raw.cursor() as cur:
            cur.copy_expert(
                f"COPY {table} ({columns}) FROM STDIN WITH (FORMAT csv, NULL '')",
                buffer,
            )
        raw.commit()
    finally:
        raw.close()
    logging.info("COPY loaded %d rows into %s", len(frame), table)
    return len(frame)


def build_geometry() -> None:
    """Create Point geometry for stops and LineString geometry per shape_id."""
    with ENGINE.begin() as conn:
        conn.execute(text(
            "ALTER TABLE stops "
            "ADD COLUMN IF NOT EXISTS geom geometry(Point, 4326)"
        ))
        conn.execute(text(
            "UPDATE stops "
            "SET geom = ST_SetSRID(ST_MakePoint(stop_lon, stop_lat), 4326)"
        ))
        # Aggregate shape points into a LineString, ordered by sequence
        conn.execute(text("""
            INSERT INTO shape_geom (shape_id, geom)
            SELECT shape_id,
                   ST_MakeLine(
                       ST_SetSRID(ST_MakePoint(shape_pt_lon, shape_pt_lat), 4326)
                       ORDER BY shape_pt_sequence
                   )
            FROM shapes
            GROUP BY shape_id
            ON CONFLICT (shape_id) DO UPDATE SET geom = EXCLUDED.geom
        """))
    logging.info("Geometry built for stops and shapes")


def add_indexes_and_keys() -> None:
    with ENGINE.begin() as conn:
        conn.execute(text(
            "CREATE INDEX IF NOT EXISTS ix_stops_geom "
            "ON stops USING GIST (geom)"
        ))
        conn.execute(text(
            "CREATE INDEX IF NOT EXISTS ix_stop_times_stop_id "
            "ON stop_times (stop_id)"
        ))
        # Referential foreign keys between the relational tables
        conn.execute(text("""
            ALTER TABLE stop_times
            ADD CONSTRAINT fk_stop_times_stop
            FOREIGN KEY (stop_id) REFERENCES stops (stop_id) NOT VALID
        """))
        conn.execute(text("""
            ALTER TABLE stop_times
            ADD CONSTRAINT fk_stop_times_trip
            FOREIGN KEY (trip_id) REFERENCES trips (trip_id) NOT VALID
        """))
    logging.info("GiST index and foreign keys added")


def load_feed(feed_dir: str) -> None:
    feed_path = Path(feed_dir)
    create_schema()
    for table in ("stops", "trips", "stop_times", "shapes"):
        copy_table(feed_path, table)
    build_geometry()
    add_indexes_and_keys()
    logging.info("GTFS feed loaded into PostGIS from %s", feed_dir)


# --- Entry point ---
# load_feed("/data/gtfs/mta")

from sqlalchemy import text

with ENGINE.connect() as conn:
    # Every stop must have a non-null, valid Point geometry
    bad_geom = conn.execute(text(
        "SELECT count(*) FROM stops "
        "WHERE geom IS NULL OR NOT ST_IsValid(geom)"
    )).scalar_one()
    assert bad_geom == 0, f"{bad_geom} stops have missing or invalid geometry"

    # SRID must be 4326 on both geometry columns
    stop_srid = conn.execute(text(
        "SELECT DISTINCT ST_SRID(geom) FROM stops"
    )).scalar_one()
    assert stop_srid == 4326, f"stops.geom SRID is {stop_srid}, expected 4326"

    # No stop_times row may reference a stop_id absent from stops
    orphans = conn.execute(text("""
        SELECT count(*) FROM stop_times st
        LEFT JOIN stops s ON st.stop_id = s.stop_id
        WHERE s.stop_id IS NULL
    """)).scalar_one()
    assert orphans == 0, f"{orphans} stop_times reference missing stops"

    # A spatial query should now use the GiST index
    nearby = conn.execute(text("""
        SELECT count(*) FROM stops
        WHERE ST_DWithin(
            geom::geography,
            ST_SetSRID(ST_MakePoint(-73.9857, 40.7484), 4326)::geography,
            500
        )
    """)).scalar_one()
    print(f"Stops within 500 m of the test point: {nearby}")

print("PostGIS load verified — geometry, SRID, and foreign keys all pass")