swarm-house/simulator/explorer/server.py

"""Data-plane explorer: a web view over the Hive-partitioned Parquet lake.

Serves three things:
  /            single-page UI (partition tree + read-only SQL console)
  /api/tree    partition hierarchy with file counts and bytes, live
  /api/query   gated read-only DuckDB SQL, same statement rules as the
               peer query channel (SELECT/WITH only, single statement)

The point is doctrinal, not just convenient: the explorer reuses the exact
read-only SQL contract that drones expose to each other, so "looking at the
data plane" on the bench exercises the same path a peer would use in flight.
"""

from __future__ import annotations

import json
import os
import re
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

import duckdb

DATA_DIR = Path(os.environ.get("DATA_DIR", "./data"))
PORT = int(os.environ.get("EXPLORER_PORT", "8088"))
ROW_LIMIT = int(os.environ.get("ROW_LIMIT", "500"))
STATIC_DIR = Path(__file__).parent

# Same spirit as the forced-command gate on a real drone: one statement,
# must be a read, no statement that could write, configure, or reach out.
_ALLOWED_START = re.compile(r"^\s*(SELECT|WITH|DESCRIBE|SUMMARIZE|SHOW)\b", re.IGNORECASE)
_FORBIDDEN = re.compile(
    r"\b(INSERT|UPDATE|DELETE|CREATE|DROP|ALTER|ATTACH|DETACH|COPY|EXPORT"
    r"|INSTALL|LOAD|SET|RESET|PRAGMA|CALL|IMPORT|VACUUM|CHECKPOINT)\b",
    re.IGNORECASE,
)


def gate(sql: str) -> str | None:
    """Return a rejection reason, or None if the statement passes."""
    stripped = re.sub(r"--[^\n]*|/\*.*?\*/", " ", sql, flags=re.DOTALL).strip().rstrip(";")
    if not stripped:
        return "empty statement"
    if ";" in stripped:
        return "multiple statements are not allowed"
    if not _ALLOWED_START.match(stripped):
        return "only SELECT / WITH / DESCRIBE / SUMMARIZE / SHOW are allowed"
    if _FORBIDDEN.search(stripped):
        return "statement contains a forbidden keyword"
    return None


def connect() -> duckdb.DuckDBPyConnection:
    """Fresh connection with the three datasets pre-registered as views."""
    con = duckdb.connect()
    for ds in ("telemetry", "detections", "state"):
        pattern = f"{DATA_DIR}/dataset={ds}/**/*.parquet"
        try:
            con.execute(
                f"CREATE VIEW {ds} AS SELECT * FROM read_parquet("
                f"'{pattern}', hive_partitioning=true, union_by_name=true)"
            )
        except duckdb.Error:
            pass  # dataset not written yet; view simply won't exist
    return con


def tree() -> dict:
    """Partition hierarchy: dataset -> flight -> drone -> leafs, with sizes."""
    root: dict = {}
    total_bytes = 0
    total_files = 0
    for f in sorted(DATA_DIR.rglob("*.parquet")):
        rel = f.relative_to(DATA_DIR)
        size = f.stat().st_size
        total_bytes += size
        total_files += 1
        node = root
        for part in rel.parts[:-1]:
            node = node.setdefault("children", {}).setdefault(part, {})
        leaf = node.setdefault("children", {}).setdefault(rel.parts[-1], {})
        leaf["bytes"] = size
        # roll sizes up the tree
        node = root
        node["bytes"] = node.get("bytes", 0) + size
        node["files"] = node.get("files", 0) + 1
        for part in rel.parts[:-1]:
            node = node["children"][part]
            node["bytes"] = node.get("bytes", 0) + size
            node["files"] = node.get("files", 0) + 1
    return {"tree": root, "total_bytes": total_bytes, "total_files": total_files}


def run_query(sql: str) -> dict:
    reason = gate(sql)
    if reason:
        return {"error": f"rejected by read-only gate: {reason}"}
    con = connect()
    try:
        cur = con.sql(sql)
        columns = cur.columns
        rows = cur.fetchmany(ROW_LIMIT)
        return {
            "columns": columns,
            "rows": [[repr(v) if isinstance(v, bytes) else v for v in row] for row in rows],
            "truncated": len(rows) == ROW_LIMIT,
        }
    except duckdb.Error as exc:
        return {"error": str(exc)}
    finally:
        con.close()


class Handler(BaseHTTPRequestHandler):
    def _send(self, code: int, body: bytes, ctype: str) -> None:
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _json(self, payload: dict, code: int = 200) -> None:
        self._send(code, json.dumps(payload, default=str).encode(), "application/json")

    def do_GET(self) -> None:  # noqa: N802 — http.server API
        if self.path in ("/", "/index.html"):
            self._send(200, (STATIC_DIR / "index.html").read_bytes(), "text/html; charset=utf-8")
        elif self.path == "/api/tree":
            self._json(tree())
        else:
            self._send(404, b"not found", "text/plain")

    def do_POST(self) -> None:  # noqa: N802 — http.server API
        if self.path != "/api/query":
            self._send(404, b"not found", "text/plain")
            return
        length = int(self.headers.get("Content-Length", "0"))
        try:
            sql = json.loads(self.rfile.read(length)).get("sql", "")
        except (json.JSONDecodeError, AttributeError):
            self._json({"error": "invalid request body"}, 400)
            return
        self._json(run_query(sql))

    def log_message(self, *_args: object) -> None:
        pass


if __name__ == "__main__":
    print(f"data-plane explorer on :{PORT}, reading {DATA_DIR}")
    ThreadingHTTPServer(("", PORT), Handler).serve_forever()