atc-agents/api/hdfs_api.py

"""HDFS (WebHDFS) file browser API for Command Center — mirrors the S3 browser."""

from __future__ import annotations

import os
from datetime import datetime, timezone
from typing import Any
from urllib.parse import quote

import httpx
from fastapi import APIRouter, Query
from fastapi.responses import JSONResponse, StreamingResponse

HDFS_NN_URL = os.getenv("HDFS_NN_URL", "http://10.0.21.61:9870").rstrip("/")
HDFS_USER = os.getenv("HDFS_USER", "hdfs")
WEBHDFS = f"{HDFS_NN_URL}/webhdfs/v1"

router = APIRouter(prefix="/api/storage/hdfs", tags=["hdfs"])


def _human_size(n: int | float | None) -> str:
    if not n:
        return "0 B"
    n = float(n)
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if n < 1024:
            return f"{n:.0f} {unit}" if unit == "B" else f"{n:.1f} {unit}"
        n /= 1024
    return f"{n:.1f} PB"


def _perm_to_rwx(perm: str, is_dir: bool) -> str:
    try:
        bits = int(str(perm), 8)
    except ValueError:
        return ("d" if is_dir else "-") + "---------"
    chars = ["d" if is_dir else "-"]
    for shift in (6, 3, 0):
        triplet = (bits >> shift) & 0o7
        chars.append("r" if triplet & 4 else "-")
        chars.append("w" if triplet & 2 else "-")
        chars.append("x" if triplet & 1 else "-")
    return "".join(chars)


def _ms_to_iso(ms: int | None) -> str:
    if not ms:
        return ""
    try:
        return datetime.fromtimestamp(ms / 1000, tz=timezone.utc).isoformat()
    except Exception:
        return ""


def _norm(path: str) -> str:
    p = (path or "/").strip()
    if not p.startswith("/"):
        p = "/" + p
    if len(p) > 1 and p.endswith("/"):
        p = p.rstrip("/")
    return p or "/"


def _encode(path: str) -> str:
    # keep the slash structure, encode each segment
    return "/".join(quote(seg, safe="") for seg in path.split("/"))


@router.get("/health")
async def hdfs_health() -> dict[str, Any]:
    try:
        async with httpx.AsyncClient(timeout=6.0) as client:
            jmx = f"{HDFS_NN_URL}/jmx?qry=Hadoop:service=NameNode,name=FSNamesystemState"
            r = await client.get(jmx)
            beans = (r.json().get("beans") or [{}])[0]
            cap = beans.get("CapacityTotal") or 0
            used = beans.get("CapacityUsed") or 0
            return {
                "ok": True,
                "namenode": HDFS_NN_URL,
                "user": HDFS_USER,
                "live_datanodes": beans.get("NumLiveDataNodes"),
                "dead_datanodes": beans.get("NumDeadDataNodes"),
                "capacity_total_human": _human_size(cap),
                "capacity_used_human": _human_size(used),
                "capacity_used_pct": round(used / cap * 100, 1) if cap else 0,
            }
    except Exception as exc:
        return {"ok": False, "namenode": HDFS_NN_URL, "error": str(exc)}


@router.get("/list")
async def hdfs_list(path: str = Query("/")) -> JSONResponse:
    p = _norm(path)
    url = f"{WEBHDFS}{_encode(p)}"
    try:
        async with httpx.AsyncClient(timeout=15.0) as client:
            r = await client.get(url, params={"op": "LISTSTATUS", "user.name": HDFS_USER})
        if r.status_code >= 400:
            try:
                msg = r.json().get("RemoteException", {}).get("message", r.text[:200])
            except Exception:
                msg = r.text[:200]
            return JSONResponse({"ok": False, "path": p, "error": msg}, status_code=200)
        statuses = (r.json().get("FileStatuses") or {}).get("FileStatus") or []
    except Exception as exc:
        return JSONResponse({"ok": False, "path": p, "error": str(exc)}, status_code=200)

    folders: list[dict[str, Any]] = []
    files: list[dict[str, Any]] = []
    for s in statuses:
        name = s.get("pathSuffix") or ""
        is_dir = s.get("type") == "DIRECTORY"
        full = (p.rstrip("/") + "/" + name) if name else p
        entry = {
            "name": name,
            "path": full,
            "type": "directory" if is_dir else "file",
            "size": s.get("length", 0),
            "size_human": _human_size(s.get("length", 0)),
            "modified": _ms_to_iso(s.get("modificationTime")),
            "owner": s.get("owner"),
            "group": s.get("group"),
            "perms": _perm_to_rwx(s.get("permission", "000"), is_dir),
            "replication": s.get("replication"),
        }
        (folders if is_dir else files).append(entry)

    folders.sort(key=lambda e: e["name"].lower())
    files.sort(key=lambda e: e["name"].lower())
    return JSONResponse({"ok": True, "path": p, "folders": folders, "files": files})


@router.get("/summary")
async def hdfs_summary(path: str = Query("/")) -> dict[str, Any]:
    p = _norm(path)
    url = f"{WEBHDFS}{_encode(p)}"
    try:
        async with httpx.AsyncClient(timeout=15.0) as client:
            r = await client.get(url, params={"op": "GETCONTENTSUMMARY", "user.name": HDFS_USER})
        cs = r.json().get("ContentSummary") or {}
        return {
            "ok": True,
            "path": p,
            "directories": cs.get("directoryCount"),
            "files": cs.get("fileCount"),
            "length": cs.get("length"),
            "length_human": _human_size(cs.get("length")),
            "space_consumed": cs.get("spaceConsumed"),
            "space_consumed_human": _human_size(cs.get("spaceConsumed")),
        }
    except Exception as exc:
        return {"ok": False, "path": p, "error": str(exc)}


@router.get("/download")
async def hdfs_download(path: str = Query(...)):
    p = _norm(path)
    url = f"{WEBHDFS}{_encode(p)}"
    fname = p.rsplit("/", 1)[-1] or "download"

    async def stream():
        async with httpx.AsyncClient(timeout=None, follow_redirects=True, verify=False) as client:
            async with client.stream("GET", url, params={"op": "OPEN", "user.name": HDFS_USER}) as resp:
                if resp.status_code >= 400:
                    return
                async for chunk in resp.aiter_bytes(chunk_size=65536):
                    yield chunk

    return StreamingResponse(
        stream(),
        media_type="application/octet-stream",
        headers={"Content-Disposition": f'attachment; filename="{fname}"'},
    )


@router.get("/preview")
async def hdfs_preview(path: str = Query(...), max_bytes: int = Query(65536)) -> dict[str, Any]:
    """Return the first chunk of a file as text (for quick inspection)."""
    p = _norm(path)
    url = f"{WEBHDFS}{_encode(p)}"
    cap = max(1024, min(max_bytes, 512 * 1024))
    try:
        async with httpx.AsyncClient(timeout=20.0, follow_redirects=True, verify=False) as client:
            r = await client.get(url, params={"op": "OPEN", "user.name": HDFS_USER, "length": cap})
        if r.status_code >= 400:
            return {"ok": False, "path": p, "error": f"HTTP {r.status_code}"}
        raw = r.content[:cap]
        try:
            text = raw.decode("utf-8")
            binary = False
        except UnicodeDecodeError:
            text = raw[:4096].decode("latin-1", errors="replace")
            binary = True
        return {"ok": True, "path": p, "binary": binary, "truncated": len(r.content) >= cap, "text": text}
    except Exception as exc:
        return {"ok": False, "path": p, "error": str(exc)}