Generating Valid PostGIS Queries from Natural Language

import re
from typing import Set
from pydantic import BaseModel, field_validator

class CoordinateBounds(BaseModel):
    """Enforces WGS84 and common projected coordinate limits."""
    min_lon: float = -180.0
    max_lon: float = 180.0
    min_lat: float = -90.0
    max_lat: float = 90.0

class SpatialQueryContract(BaseModel):
    """Strict schema contract for generated PostGIS queries."""
    table_name: str
    geometry_column: str
    target_srid: int
    raw_query: str
    allowed_srids: Set[int] = frozenset({4326, 3857, 26910, 4269})

    @field_validator('target_srid')
    @classmethod
    def validate_srid(cls, v: int) -> int:
        allowed = {4326, 3857, 26910, 4269}
        if v not in allowed:
            raise ValueError(f"Unsupported SRID {v}. Must be one of {allowed}.")
        return v

    @field_validator('raw_query')
    @classmethod
    def validate_coordinates_and_syntax(cls, v: str) -> str:
        # Extract coordinates from ST_MakePoint inline literals
        coord_pattern = r"ST_MakePoint\(\s*([\d\.\-]+)\s*,\s*([\d\.\-]+)"
        matches = re.findall(coord_pattern, v)

        bounds = CoordinateBounds()
        for lon_str, lat_str in matches:
            lon, lat = float(lon_str), float(lat_str)
            if not (bounds.min_lon <= lon <= bounds.max_lon):
                raise ValueError(f"Longitude {lon} violates WGS84 bounds [-180, 180]")
            if not (bounds.min_lat <= lat <= bounds.max_lat):
                raise ValueError(f"Latitude {lat} violates WGS84 bounds [-90, 90]")

        if "DROP" in v.upper() or "TRUNCATE" in v.upper():
            raise ValueError("DDL/DML destructive operations are strictly prohibited.")

        return v

# Next Steps for Pipeline Integration:
# 1. Register SpatialQueryContract in your API gateway's request validation middleware.
# 2. Cache allowed_srids dynamically by querying `SELECT srid FROM spatial_ref_sys` at startup.
# 3. Integrate with OpenTelemetry to emit `validation_failed` spans when ValidationError is raised.

import sqlparse
from typing import Dict, Any

ALLOWED_SPATIAL_FUNCTIONS = {
    "st_dwithin", "st_intersects", "st_transform",
    "st_buffer", "st_centroid", "st_area", "st_distance", "st_union"
}

# Function pairs that risk GEOS exceptions when chained without validity checks
TOPOLOGY_RISK_CHAINS = {
    frozenset({"st_buffer", "st_union"}),
    frozenset({"st_intersection", "st_difference"}),
}

def validate_ast_and_topology(query: str) -> Dict[str, Any]:
    parsed = sqlparse.parse(query)
    if not parsed:
        raise ValueError("Empty or malformed SQL payload")

    stmt = parsed[0]
    found_functions = set()

    for token in stmt.flatten():
        if token.ttype in (sqlparse.tokens.Name, sqlparse.tokens.Keyword):
            func_name = token.value.lower()
            if func_name in ALLOWED_SPATIAL_FUNCTIONS:
                found_functions.add(func_name)

    # Enforce index-friendly bounding box pre-filter (&&)
    needs_bbox = bool(found_functions & {"st_dwithin", "st_intersects"})
    if needs_bbox and "&&" not in query:
        raise ValueError(
            "Missing bounding-box pre-filter (&&). "
            "Queries using ST_Intersects or ST_DWithin will cause full-table scans without it."
        )

    # Detect high-risk topology chains
    for risk_pair in TOPOLOGY_RISK_CHAINS:
        if risk_pair.issubset(found_functions):
            funcs = " + ".join(sorted(risk_pair))
            raise RuntimeError(
                f"High-risk topology chain detected: {funcs}. "
                "Route to async worker with ST_IsValid pre-check enabled."
            )

    return {
        "functions_used": list(found_functions),
        "index_safe": "&&" in query or not needs_bbox,
        "safe_to_execute_sync": True
    }

# Next Steps for Pipeline Integration:
# 1. Hook this validator into your LLM output router before query dispatch.
# 2. Use `sqlparse.format(query, reindent=True)` for audit logging in compliance pipelines.
# 3. Implement a fallback ST_IsValid(geom) pre-check for any query containing topology chains.

import asyncpg
import asyncio
from typing import Optional, Dict, Any

class SpatialExecutionError(Exception):
    """Custom exception for structured spatial API error mapping."""
    def __init__(self, message: str, error_code: str, query: str, details: Optional[Dict] = None):
        super().__init__(message)
        self.error_code = error_code
        self.query = query
        self.details = details or {}

async def execute_validated_query(
    pool: asyncpg.Pool,
    query: str,
    timeout: float = 5.0,
    max_rows: int = 10000
) -> Dict[str, Any]:
    async with pool.acquire() as conn:
        try:
            await conn.execute(f"SET statement_timeout = '{int(timeout * 1000)}ms'")
            await conn.execute("SET work_mem = '256MB'")

            result = await conn.fetch(query)
            if len(result) > max_rows:
                raise SpatialExecutionError(
                    f"Result set exceeds {max_rows} row limit",
                    "ERR_ROW_LIMIT_EXCEEDED",
                    query
                )

            return {"status": "success", "rows": len(result), "data": [dict(r) for r in result]}

        except asyncpg.exceptions.InvalidTextRepresentationError as e:
            raise SpatialExecutionError("Coordinate parsing failed", "ERR_COORD_PARSE", query, {"original": str(e)}) from e
        except asyncpg.exceptions.UndefinedTableError as e:
            raise SpatialExecutionError("Target table or column missing", "ERR_SCHEMA_MISSING", query, {"original": str(e)}) from e
        except asyncpg.exceptions.QueryCanceledError as e:
            raise SpatialExecutionError("Query exceeded timeout or memory limit", "ERR_TIMEOUT", query, {"original": str(e)}) from e
        except asyncpg.exceptions.DataError as e:
            if "TopologyException" in str(e) or "GEOS" in str(e):
                raise SpatialExecutionError("Invalid geometry topology encountered", "ERR_TOPOLOGY_INVALID", query, {"original": str(e)}) from e
            raise SpatialExecutionError(f"Spatial data constraint violation: {str(e)}", "ERR_DATA_CONSTRAINT", query) from e
        except Exception as e:
            raise SpatialExecutionError(f"Unhandled spatial execution failure: {str(e)}", "ERR_UNKNOWN", query) from e

# Next Steps for Pipeline Integration:
# 1. Deploy this executor behind a FastAPI/Starlette endpoint with connection pooling (min=5, max=50).
# 2. Map SpatialExecutionError to HTTP 400/422/503 responses in your API error handler.
# 3. Implement exponential backoff with jitter for ERR_TIMEOUT retries in batch processing jobs.