atc-agents/api/lake_meta.py

"""Shared lakehouse metadata + Trino helpers.

Single source of truth for the business datasets the governance / data-quality /
lineage / observability features operate on, plus a small synchronous Trino
client. Imported by dq_monitor.py, observability.py, lineage.py and
catalog_governance.py so every feature reasons about the exact same tables that
the rest of the Command Center (pii_catalog, etl_offload, trino_federated)
already exposes.
"""

from __future__ import annotations

import os
from typing import Any

import httpx

TRINO_URL = os.getenv("TRINO_URL", "http://10.0.21.50:8089").rstrip("/")
TRINO_USER = os.getenv("TRINO_USER", "mo")

# Canonical business datasets, aligned with pii_catalog.DATASETS / etl_offload.
# fqtn  = fully-qualified Trino table name (catalog.schema.table)
# om_fqn = OpenMetadata FQN best-effort (governance falls back to native store)
DATASETS: list[dict[str, Any]] = [
    {"key": "orders", "label": "Sales orders", "engine": "PostgreSQL", "color": "#fbbf24",
     "catalog": "postgres_sales", "schema": "public", "table": "sales_orders",
     "fqtn": "postgres_sales.public.sales_orders", "pii_key": "postgres",
     "om_fqn": "atc_trino.postgres_sales.public.sales_orders",
     "key_col": "order_id", "ts_col": "order_ts", "amount_col": "amount",
     "native_count": ("postgres", "public.sales_orders"),
     "domain": "Sales", "topic": "atc.public.sales_orders"},
    {"key": "hr", "label": "HR events", "engine": "MySQL", "color": "#60a5fa",
     "catalog": "mysql_hr", "schema": "hr", "table": "employee_events",
     "fqtn": "mysql_hr.hr.employee_events", "pii_key": "mysql",
     "om_fqn": "atc_trino.mysql_hr.hr.employee_events",
     "key_col": "event_id", "ts_col": "event_ts", "amount_col": "salary_change",
     "native_count": ("mysql", "hr.employee_events"),
     "domain": "People", "topic": "atc.hr.employee_events"},
    {"key": "supply", "label": "Supply chain events", "engine": "MongoDB", "color": "#a78bfa",
     "catalog": "mongodb_supplychain", "schema": "supplychain", "table": "events",
     "fqtn": "mongodb_supplychain.supplychain.events", "pii_key": "mongodb",
     "om_fqn": "atc_trino.mongodb_supplychain.supplychain.events",
     "key_col": "event_id", "ts_col": "ts", "amount_col": "amount",
     "native_count": ("mongodb", "supplychain.events"),
     "domain": "Supply Chain", "topic": "atc.supplychain.events"},
    {"key": "telemetry", "label": "Device telemetry", "engine": "Cassandra", "color": "#22d3ee",
     "catalog": "cassandra_telemetry", "schema": "telemetry", "table": "device_metrics",
     "fqtn": "cassandra_telemetry.telemetry.device_metrics", "pii_key": "cassandra",
     "om_fqn": "atc_trino.cassandra_telemetry.telemetry.device_metrics",
     "key_col": "device_id", "ts_col": "metric_ts", "amount_col": "metric_value",
     "unique_key": False,
     "domain": "IoT", "topic": "atc.telemetry.device_metrics"},
    {"key": "curated", "label": "Curated masked (Iceberg)", "engine": "Iceberg / Trino", "color": "#34d399",
     "catalog": "iceberg", "schema": "curated_masked", "table": "sales_orders_masked",
     "fqtn": "iceberg.curated_masked.sales_orders_masked", "pii_key": "curated",
     "om_fqn": "atc_trino.iceberg.curated_masked.sales_orders_masked",
     "key_col": "order_id", "ts_col": "order_ts", "amount_col": "amount",
     "domain": "Sales", "curated": True, "topic": None},
    {"key": "hadoop", "label": "Historical sales (HDFS)", "engine": "Iceberg / HDFS", "color": "#f472b6",
     "catalog": "iceberg", "schema": "hadoop", "table": "historical_sales_hdfs",
     "fqtn": "iceberg.hadoop.historical_sales_hdfs", "pii_key": "hadoop",
     "om_fqn": "atc_trino.iceberg.hadoop.historical_sales_hdfs",
     "key_col": "order_id", "ts_col": "order_ts", "amount_col": "amount",
     "domain": "Sales", "curated": True, "topic": None},
]

DATASET_BY_KEY = {d["key"]: d for d in DATASETS}

# Heuristics for picking the freshness / key column when discovering schema.
_TS_HINTS = ("_ts", "ts", "updated_at", "created_at", "event_time", "modified", "time")
_KEY_HINTS = ("_id", "id", "uuid", "key", "pk")


def trino(sql: str, timeout: float = 25.0) -> tuple[list[str], list[list[Any]]]:
    """Run a Trino statement, following nextUri pages. Returns (columns, rows).

    `timeout` is both the per-request timeout AND an overall wall-clock deadline,
    so a long full-scan (e.g. count(*) on a huge Cassandra table) is aborted and
    cancelled instead of looping over nextUri pages for minutes and blocking the
    caller's thread.
    """
    import time as _t
    cols: list[str] = []
    rows: list[list[Any]] = []
    deadline = _t.monotonic() + timeout
    with httpx.Client(timeout=min(timeout, 15.0)) as client:
        d = client.post(f"{TRINO_URL}/v1/statement", content=sql.encode(),
                        headers={"X-Trino-User": TRINO_USER}).json()
        while True:
            if d.get("error"):
                raise RuntimeError(d["error"].get("message", "trino error"))
            c = d.get("columns")
            if c and not cols:
                cols = [x["name"] for x in c]
            rows += d.get("data") or []
            nxt = d.get("nextUri")
            if not nxt:
                break
            if _t.monotonic() > deadline:
                try:
                    client.delete(nxt)  # cancel the running query server-side
                except Exception:
                    pass
                raise TimeoutError(f"trino query exceeded {timeout}s deadline")
            d = client.get(nxt).json()
    return cols, rows


def trino_scalar(sql: str, timeout: float = 25.0) -> Any:
    _c, rows = trino(sql, timeout=timeout)
    if rows and rows[0]:
        return rows[0][0]
    return None


def discover_columns(ds: dict[str, Any], timeout: float = 12.0) -> list[dict[str, str]]:
    """[{name, type}] for a dataset's table via information_schema."""
    sql = (f"SELECT column_name, data_type FROM {ds['catalog']}.information_schema.columns "
           f"WHERE table_name = '{ds['table']}'")
    if ds.get("schema"):
        sql += f" AND table_schema = '{ds['schema']}'"
    try:
        _c, rows = trino(sql, timeout=timeout)
        return [{"name": r[0], "type": r[1]} for r in rows]
    except Exception:
        return []


def pick_ts_col(columns: list[dict[str, str]], default: str | None = None) -> str | None:
    names = [c["name"] for c in columns]
    for c in columns:
        if "timestamp" in (c.get("type") or "").lower() or "date" in (c.get("type") or "").lower():
            return c["name"]
    for n in names:
        if any(h in n.lower() for h in _TS_HINTS):
            return n
    return default if default in names else None


def pick_key_col(columns: list[dict[str, str]], default: str | None = None) -> str | None:
    names = [c["name"] for c in columns]
    if default in names:
        return default
    for n in names:
        if any(n.lower() == h or n.lower().endswith(h) for h in _KEY_HINTS):
            return n
    return names[0] if names else None