Designing a GTFS Schema for Postgres

-- gtfs_schema.sql — one schema per load, swapped in atomically.

CREATE SCHEMA IF NOT EXISTS gtfs_staging;
SET search_path TO gtfs_staging;

-- Every table carries the feed version, so two publications can coexist and a
-- reload never has to destroy the version that is currently being served.
CREATE TABLE feed_version (
    feed_version    text PRIMARY KEY,       -- the content checksum of the archive
    fetched_at      timestamptz NOT NULL,
    feed_start_date date,
    feed_end_date   date,
    agency_timezone text NOT NULL
);

CREATE TABLE routes (
    feed_version     text NOT NULL REFERENCES feed_version ON DELETE CASCADE,
    route_id         text NOT NULL,
    agency_id        text,
    route_short_name text,
    route_long_name  text,
    route_type       smallint NOT NULL,
    route_color      char(6),               -- hex; text, so 00FF00 survives
    PRIMARY KEY (feed_version, route_id)
);

CREATE TABLE calendar_service (
    feed_version text NOT NULL REFERENCES feed_version ON DELETE CASCADE,
    service_id   text NOT NULL,
    service_date date NOT NULL,             -- the EXPANDED calendar, one row per date
    PRIMARY KEY (feed_version, service_id, service_date)
);

CREATE TABLE trips (
    feed_version text NOT NULL REFERENCES feed_version ON DELETE CASCADE,
    trip_id      text NOT NULL,
    route_id     text NOT NULL,
    service_id   text NOT NULL,
    shape_id     text,
    block_id     text,
    direction_id smallint,
    trip_headsign text,
    PRIMARY KEY (feed_version, trip_id),
    CONSTRAINT trips_route_fk FOREIGN KEY (feed_version, route_id)
        REFERENCES routes DEFERRABLE INITIALLY DEFERRED
);

CREATE TABLE stops (
    feed_version   text NOT NULL REFERENCES feed_version ON DELETE CASCADE,
    stop_id        text NOT NULL,
    stop_name      text,
    stop_code      text,
    location_type  smallint NOT NULL DEFAULT 0,
    parent_station text,
    zone_id        text,
    geom           geometry(Point, 4326),   -- as published; project at query time
    PRIMARY KEY (feed_version, stop_id)
);

CREATE TABLE stop_times (
    feed_version        text NOT NULL REFERENCES feed_version ON DELETE CASCADE,
    trip_id             text NOT NULL,
    stop_sequence       integer NOT NULL,
    stop_id             text NOT NULL,
    -- Seconds since the service day began. GTFS values exceed 24 hours, which
    -- the time type rejects outright; integers also index and sort correctly
    -- across midnight, which is the whole point.
    arrival_s           integer,
    departure_s         integer,
    pickup_type         smallint NOT NULL DEFAULT 0,
    drop_off_type       smallint NOT NULL DEFAULT 0,
    shape_dist_traveled double precision,
    timepoint           smallint,
    PRIMARY KEY (feed_version, trip_id, stop_sequence),
    CONSTRAINT stop_times_trip_fk FOREIGN KEY (feed_version, trip_id)
        REFERENCES trips DEFERRABLE INITIALLY DEFERRED,
    CONSTRAINT stop_times_stop_fk FOREIGN KEY (feed_version, stop_id)
        REFERENCES stops DEFERRABLE INITIALLY DEFERRED,
    CONSTRAINT stop_times_ordered CHECK (
        arrival_s IS NULL OR departure_s IS NULL OR departure_s >= arrival_s)
);

"""Create the schema, load into it, index it, and swap it in atomically."""
from __future__ import annotations

import logging
from pathlib import Path

import psycopg

log = logging.getLogger("gtfs.postgres")

# Built AFTER the data is in: maintaining them during a bulk COPY roughly
# doubles the load time and buys nothing, because nothing queries mid-load.
POST_LOAD_INDEXES = (
    # The query every departure board runs.
    "CREATE INDEX stop_times_by_stop ON stop_times (feed_version, stop_id, departure_s)",
    # Trip reconstruction, in stop order.
    "CREATE INDEX stop_times_by_trip ON stop_times (feed_version, trip_id, stop_sequence)",
    # 'What runs today' — the calendar lookup behind almost every query.
    "CREATE INDEX calendar_by_date ON calendar_service (feed_version, service_date)",
    "CREATE INDEX trips_by_route ON trips (feed_version, route_id)",
    "CREATE INDEX trips_by_service ON trips (feed_version, service_id)",
    "CREATE INDEX stops_geom ON stops USING gist (geom)",
    "CREATE INDEX stops_by_parent ON stops (feed_version, parent_station)"
    " WHERE parent_station IS NOT NULL",
)


def load_feed(dsn: str, schema_sql: Path, copy_tables: dict[str, str],
              feed_version: str) -> None:
    """copy_tables: table name -> a CSV buffer path, already normalised."""
    with psycopg.connect(dsn) as conn:
        with conn.cursor() as cur:
            cur.execute("DROP SCHEMA IF EXISTS gtfs_staging CASCADE")
            cur.execute(schema_sql.read_text())
        conn.commit()

        with conn.cursor() as cur:
            # Deferred constraints mean the tables may load in ANY order, so this
            # loop can become a thread pool without a dependency graph.
            cur.execute("SET CONSTRAINTS ALL DEFERRED")
            for table, csv_path in copy_tables.items():
                with open(csv_path, "rb") as fh, cur.copy(
                        f"COPY gtfs_staging.{table} FROM STDIN WITH (FORMAT csv, HEADER)"
                ) as copy:
                    while chunk := fh.read(1 << 20):
                        copy.write(chunk)
                log.info("copied %s", table)
        conn.commit()          # every deferred foreign key is checked here

        with conn.cursor() as cur:
            for statement in POST_LOAD_INDEXES:
                cur.execute(statement.replace(" ON ", " ON gtfs_staging.", 1))
            cur.execute("ANALYZE")
        conn.commit()
        log.info("staging schema ready for feed_version %s", feed_version)


def promote(dsn: str, keep_previous: bool = True) -> None:
    """Swap staging into place in ONE transaction; readers never see a partial feed."""
    with psycopg.connect(dsn) as conn, conn.cursor() as cur:
        cur.execute("DROP SCHEMA IF EXISTS gtfs_previous CASCADE")
        cur.execute("ALTER SCHEMA gtfs RENAME TO gtfs_previous")
        cur.execute("ALTER SCHEMA gtfs_staging RENAME TO gtfs")
        if not keep_previous:
            cur.execute("DROP SCHEMA gtfs_previous CASCADE")
        conn.commit()
        log.info("promoted staging to live%s",
                 "; previous kept for rollback" if keep_previous else "")

VERIFY_SQL = """
-- Every trip has calls, and every call belongs to a trip.
SELECT 'trips with no stop_times' AS check, count(*) AS n
FROM gtfs_staging.trips t
WHERE NOT EXISTS (SELECT 1 FROM gtfs_staging.stop_times st
                  WHERE st.feed_version = t.feed_version AND st.trip_id = t.trip_id)
UNION ALL
-- Every trip's service resolves to at least one real date.
SELECT 'trips whose service never runs', count(*)
FROM gtfs_staging.trips t
WHERE NOT EXISTS (SELECT 1 FROM gtfs_staging.calendar_service c
                  WHERE c.feed_version = t.feed_version AND c.service_id = t.service_id)
UNION ALL
-- Times must not run backwards within a trip.
SELECT 'trips with non-monotonic times', count(*) FROM (
    SELECT trip_id FROM (
        SELECT trip_id, departure_s,
               lag(departure_s) OVER (PARTITION BY feed_version, trip_id
                                      ORDER BY stop_sequence) AS previous
        FROM gtfs_staging.stop_times) s
    WHERE previous IS NOT NULL AND departure_s < previous
    GROUP BY trip_id) bad
UNION ALL
SELECT 'stops with no geometry', count(*)
FROM gtfs_staging.stops WHERE geom IS NULL AND location_type IN (0, 4);
"""


def verify(dsn: str) -> None:
    with psycopg.connect(dsn) as conn, conn.cursor() as cur:
        cur.execute(VERIFY_SQL)
        for check, n in cur.fetchall():
            log.info("%-34s %d", check, n)
            assert n == 0, f"{check}: {n}"

INFO gtfs.postgres: copied routes
INFO gtfs.postgres: copied stops
INFO gtfs.postgres: copied trips
INFO gtfs.postgres: copied calendar_service
INFO gtfs.postgres: copied stop_times
INFO gtfs.postgres: staging schema ready for feed_version 9f2a1c...
INFO gtfs.postgres: trips with no stop_times           0
INFO gtfs.postgres: trips whose service never runs     0
INFO gtfs.postgres: trips with non-monotonic times     0
INFO gtfs.postgres: stops with no geometry             0
INFO gtfs.postgres: promoted staging to live; previous kept for rollback