Capturing Transformation Lineage for BigQuery GIS Jobs

from __future__ import annotations

import json
import re
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from google.cloud import bigquery

ST_PATTERN = re.compile(r"\b(ST_[A-Z_]+)\s*\(", re.IGNORECASE)


def run_and_capture_lineage(
    client: bigquery.Client,
    sql: str,
    region: str,
    lineage_dir: str | Path,
) -> dict[str, Any]:
    """Run a BigQuery GIS query and write a JSON lineage document for the job.

    Args:
        client: An authenticated BigQuery client.
        sql: The GIS SQL to execute (may contain ST_* functions).
        region: Region qualifier for INFORMATION_SCHEMA, e.g. "region-us".
        lineage_dir: Directory that receives the .json lineage document.

    Returns:
        The lineage document written to disk.
    """
    out_dir = Path(lineage_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    query_job = client.query(sql)
    query_job.result()  # block until the job completes
    job_id = query_job.job_id

    # Pull authoritative job metadata from the region-scoped JOBS view.
    meta_sql = f"""
        SELECT
          job_id,
          creation_time,
          total_bytes_processed,
          total_slot_ms,
          destination_table,
          referenced_tables
        FROM `{region}`.INFORMATION_SCHEMA.JOBS
        WHERE job_id = @job_id
    """
    cfg = bigquery.QueryJobConfig(
        query_parameters=[bigquery.ScalarQueryParameter("job_id", "STRING", job_id)]
    )
    row = next(iter(client.query(meta_sql, job_config=cfg).result()))

    def _fqtn(t: Any) -> str:
        return f"{t['project_id']}.{t['dataset_id']}.{t['table_id']}"

    inputs = [_fqtn(t) for t in (row.referenced_tables or [])]
    dest = row.destination_table
    output = _fqtn(dest) if dest else None
    spatial_ops = sorted({m.group(1).upper() for m in ST_PATTERN.finditer(sql)})

    lineage: dict[str, Any] = {
        "event": "bigquery_gis_transform",
        "job_id": job_id,
        "inputs": inputs,
        "output": output,
        "spatial_operations": spatial_ops,
        "total_bytes_processed": int(row.total_bytes_processed or 0),
        "total_slot_ms": int(row.total_slot_ms or 0),
        "job_created_at": row.creation_time.isoformat(),
        "captured_at": datetime.now(timezone.utc).isoformat(),
    }

    (out_dir / f"{job_id.replace(':', '_')}.json").write_text(
        json.dumps(lineage, indent=2), encoding="utf-8"
    )
    return lineage


if __name__ == "__main__":
    bq = bigquery.Client()
    doc = run_and_capture_lineage(
        client=bq,
        sql="""
            CREATE OR REPLACE TABLE geo.flood_parcels AS
            SELECT p.parcel_id, p.geom
            FROM geo.parcels AS p, geo.flood_zones AS f
            WHERE ST_Intersects(p.geom, f.geom)
        """,
        region="region-us",
        lineage_dir="./lineage",
    )
    print("Captured", doc["spatial_operations"], "over", doc["inputs"])

SELECT job_id, referenced_tables, destination_table
FROM `region-us`.INFORMATION_SCHEMA.JOBS
WHERE job_id = 'your_project:US.bquxjob_1a2b3c4d_00'