Tracing PostGIS Query Spans in Spatial ETL

from __future__ import annotations

import time
from contextlib import contextmanager
from typing import Any, Iterator, Optional, Sequence

import psycopg
from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode

tracer = trace.get_tracer("geospatial.postgis")

SLOW_QUERY_SECONDS = 2.0        # above this, spend a round trip on the plan


@contextmanager
def spatial_query(
    conn: psycopg.Connection,
    operation: str,                 # "intersect_parcels", not "SELECT"
    table: str,
    bbox: Optional[tuple[float, float, float, float]] = None,
    srid: int = 3857,
) -> Iterator[dict]:
    """A span named for the spatial operation, with the facts that explain it."""
    ctx: dict = {"rows": 0}
    with tracer.start_as_current_span(f"postgis.{operation}") as span:
        span.set_attribute("db.system", "postgresql")
        span.set_attribute("db.operation", operation)
        span.set_attribute("db.sql.table", table)
        span.set_attribute("geo.srid", srid)
        if bbox:
            span.set_attribute("geo.bbox", ",".join(f"{v:.3f}" for v in bbox))
            # Area is the single best predictor of how much work a spatial
            # predicate will do, and it is one multiplication.
            span.set_attribute("geo.bbox_area", (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]))

        started = time.monotonic()
        try:
            yield ctx
        except Exception as exc:
            span.record_exception(exc)
            span.set_status(Status(StatusCode.ERROR, str(exc)[:200]))
            raise
        finally:
            elapsed = time.monotonic() - started
            span.set_attribute("db.rows_returned", ctx.get("rows", 0))
            if "geometries" in ctx:
                span.set_attribute("geo.geometries_returned", ctx["geometries"])
            # The plan is expensive to fetch and useless on a fast query. Only
            # attach it where someone would actually read it.
            if elapsed > SLOW_QUERY_SECONDS and "sql" in ctx:
                span.set_attribute("db.plan", explain(conn, ctx["sql"], ctx.get("params")))


def explain(conn: psycopg.Connection, sql: str, params: Any = None) -> str:
    """A real plan with real timings — costed, not estimated."""
    with conn.cursor() as cur:
        cur.execute("EXPLAIN (ANALYZE, BUFFERS, FORMAT TEXT) " + sql, params)
        return "\n".join(row[0] for row in cur.fetchall())[:4000]


INTERSECT_SQL = """
SELECT p.parcel_id, ST_AsBinary(p.geom) AS geom
FROM   parcels p
WHERE  ST_Intersects(p.geom, ST_TileEnvelope(%(z)s, %(x)s, %(y)s))
"""


def parcels_for_tile(conn: psycopg.Connection, z: int, x: int, y: int) -> list[tuple]:
    bbox = tile_envelope(z, x, y)
    with spatial_query(conn, "intersect_parcels", "parcels", bbox=bbox) as ctx:
        ctx["sql"] = INTERSECT_SQL
        ctx["params"] = {"z": z, "x": x, "y": y}
        with conn.cursor() as cur:
            cur.execute(INTERSECT_SQL, ctx["params"])
            rows = cur.fetchall()
        ctx["rows"] = len(rows)
        ctx["geometries"] = sum(1 for r in rows if r[1] is not None)
        return rows

def test_plan_is_attached_only_when_slow(conn, exporter, monkeypatch) -> None:
    monkeypatch.setattr("mymodule.SLOW_QUERY_SECONDS", 0.0)
    parcels_for_tile(conn, 10, 550, 320)
    slow_span = exporter.get_finished_spans()[-1]
    assert "db.plan" in slow_span.attributes

    exporter.clear()
    monkeypatch.setattr("mymodule.SLOW_QUERY_SECONDS", 3600.0)
    parcels_for_tile(conn, 14, 8802, 5121)
    fast_span = exporter.get_finished_spans()[-1]
    assert "db.plan" not in fast_span.attributes


def test_span_carries_the_spatial_facts(conn, exporter) -> None:
    parcels_for_tile(conn, 14, 8802, 5121)
    span = exporter.get_finished_spans()[-1]
    assert span.name == "postgis.intersect_parcels"
    assert span.attributes["geo.srid"] == 3857
    assert span.attributes["geo.bbox_area"] > 0
    assert span.attributes["db.rows_returned"] >= 0


def test_failure_sets_error_status(conn, exporter) -> None:
    with pytest.raises(psycopg.errors.UndefinedTable):
        with spatial_query(conn, "intersect_missing", "nope") as ctx:
            with conn.cursor() as cur:
                cur.execute("SELECT * FROM nope")
    span = exporter.get_finished_spans()[-1]
    assert span.status.status_code is StatusCode.ERROR

SELECT substring(query, 1, 60) AS shape, calls,
       round(mean_exec_time::numeric, 1) AS mean_ms,
       round(total_exec_time::numeric / 1000, 1) AS total_s
FROM   pg_stat_statements
WHERE  query ILIKE '%ST_Intersects%'
ORDER  BY total_exec_time DESC
LIMIT  10;