atc-agents/api/hadoop_sql.py

"""Hadoop / HDFS / Iceberg catalog and sampling for the Data Hub."""

from __future__ import annotations

import csv
import io
import os
from typing import Any, Callable

from webhdfs_util import open_bytes

HDFS_PATHS = [
    "/data/historical/sales_orders",
    "/data/historical/sales_orders/year=2020/part-0.csv",
]

ICEBERG_SCHEMAS = ["hadoop", "curated_masked", "curated"]

HADOOP_SAMPLES: list[dict[str, str]] = [
    {"id": "hd1", "label": "Iceberg historical sales", "sql": "SELECT * FROM iceberg.hadoop.historical_sales_hdfs LIMIT 20"},
    {"id": "hd2", "label": "Row count historical", "sql": "SELECT count(*) FROM iceberg.hadoop.historical_sales_hdfs"},
    {"id": "hd3", "label": "Curated masked sample", "sql": "SELECT * FROM iceberg.curated_masked.sales_orders_masked LIMIT 15"},
    {"id": "hd4", "label": "Hive lake schemas", "sql": "SHOW SCHEMAS FROM hive"},
    {"id": "hd5", "label": "Iceberg hadoop tables", "sql": "SHOW TABLES FROM iceberg.hadoop"},
]


def catalog_hadoop(run_trino: Callable[[str, int], dict[str, Any]]) -> dict[str, Any]:
    objects: list[dict[str, Any]] = []

    for schema in ICEBERG_SCHEMAS:
        show = run_trino(f"SHOW TABLES FROM iceberg.{schema}", 500)
        if not show.get("ok"):
            continue
        for row in show.get("rows") or []:
            name = str(row[0])
            fqn = f"iceberg.{schema}.{name}"
            cnt = run_trino(f"SELECT count(*) FROM iceberg.{schema}.{name}", 1)
            count = None
            if cnt.get("ok") and cnt.get("rows"):
                try:
                    count = int(cnt["rows"][0][0])
                except (TypeError, ValueError, IndexError):
                    pass
            objects.append({
                "type": "table",
                "schema": f"iceberg.{schema}",
                "name": name,
                "fqn": fqn,
                "row_count": count,
            })

    for path in HDFS_PATHS:
        label = path.rstrip("/").rsplit("/", 1)[-1]
        objects.append({
            "type": "file",
            "schema": "hdfs",
            "name": label,
            "fqn": f"hdfs:{path}",
            "row_count": None,
        })

    return {"engine": "hadoop", "version": "HDFS + Iceberg", "objects": objects}


def sample_hdfs_csv(path: str, limit: int, offset: int, tabular: Callable[..., dict[str, Any]]) -> dict[str, Any]:
    try:
        text = open_bytes(path).decode("utf-8", errors="replace")
        reader = csv.reader(io.StringIO(text))
        all_rows = list(reader)
        if not all_rows:
            return tabular([], [], 0, row_count=0)
        columns = [c.strip() for c in all_rows[0]]
        data = all_rows[1 + offset: 1 + offset + limit]
        rows = [[cell.strip() for cell in row] for row in data]
        return tabular(columns, rows, 0, row_count=max(0, len(all_rows) - 1))
    except Exception as exc:
        return {"ok": False, "error": str(exc)[:500]}


def sample_hadoop(
    object_name: str,
    limit: int,
    offset: int,
    run_trino: Callable[[str, int], dict[str, Any]],
    tabular: Callable[..., dict[str, Any]],
) -> dict[str, Any]:
    if object_name.startswith("hdfs:"):
        return sample_hdfs_csv(object_name[5:], limit, offset, tabular)
    sql = f"SELECT * FROM {object_name} OFFSET {offset} LIMIT {limit}"
    return run_trino(sql, limit)


def table_row_count_hadoop(object_name: str, run_trino: Callable[[str, int], dict[str, Any]]) -> int | None:
    if object_name.startswith("hdfs:"):
        try:
            text = open_bytes(object_name[5:]).decode("utf-8", errors="replace")
            return max(0, sum(1 for _ in csv.reader(io.StringIO(text))) - 1)
        except Exception:
            return None
    cnt = run_trino(f"SELECT count(*) FROM {object_name}", 1)
    if cnt.get("ok") and cnt.get("rows"):
        try:
            return int(cnt["rows"][0][0])
        except (TypeError, ValueError, IndexError):
            pass
    return None


def health_hadoop(run_trino: Callable[[str, int], dict[str, Any]]) -> dict[str, Any]:
    from webhdfs_util import HDFS_NN_URL
    import httpx

    nn_ok = False
    live_dn = None
    err_parts: list[str] = []
    try:
        with httpx.Client(timeout=6.0) as client:
            r = client.get(f"{HDFS_NN_URL}/jmx?qry=Hadoop:service=NameNode,name=FSNamesystemState")
            beans = (r.json().get("beans") or [{}])[0]
            nn_ok = r.status_code < 400
            live_dn = beans.get("NumLiveDataNodes")
    except Exception as exc:
        err_parts.append(f"NN: {exc}")

    trino_ok = False
    try:
        tr = run_trino("SELECT 1", 1)
        trino_ok = bool(tr.get("ok"))
        if not trino_ok and tr.get("error"):
            err_parts.append(f"Trino: {tr['error'][:80]}")
    except Exception as exc:
        err_parts.append(f"Trino: {exc}")

    ok = nn_ok and trino_ok
    return {
        "ok": ok,
        "namenode": HDFS_NN_URL,
        "live_datanodes": live_dn,
        "trino_ok": trino_ok,
        "error": "; ".join(err_parts) if err_parts else None,
    }


def connection_info() -> dict[str, str]:
    from webhdfs_util import HDFS_NN_URL, HDFS_USER
    return {
        "namenode": HDFS_NN_URL,
        "user": HDFS_USER,
        "trino": os.getenv("TRINO_URL", "http://10.0.21.50:8089"),
        "spark": os.getenv("SPARK_UI_URL", "http://10.0.21.50:8080"),
    }