insights/products/notebooks/backend/sql_v2.py

"""Helpers for the revamped-notebooks SQLV2 run flow (Journey 1).

The backend dispatches a run to the in-sandbox kernel-server with a single HTTP
POST (mirroring Insights Desktop's agent-server). The kernel-server fetches the
node's capped result page from the data-plane endpoint (real Datastore data via
InsightsQL) and POSTs the envelope back to the token-authed callback endpoint. The
control plane (write_file / execute) is used only to deploy and launch the
kernel-server package — never per run.

The callback and data-plane tokens are stateless signed tokens for the slice;
hardening swaps them for the RS256 sandbox event-ingest JWTs used by Insights Desktop.
"""

import hmac
import time
import hashlib

from django.conf import settings
from django.core import signing

import requests
import structlog
import hanzo_insights

from insights.models.user import User

from products.notebooks.backend.kernel_package import SANDBOX_PACKAGE_NAME, kernel_package_bytes_and_hash
from products.notebooks.backend.models import KernelRuntime, Notebook, NotebookNodeRun
from products.tasks.backend.facade.sandbox import SandboxBase, get_sandbox_class_for_backend

logger = structlog.get_logger(__name__)

REVAMPED_PY_NOTEBOOKS_FLAG = "revamped-py-notebooks"

_CALLBACK_TOKEN_SALT = "notebooks.sql_v2.callback"
_CALLBACK_TOKEN_MAX_AGE_SECONDS = 3600
_DATA_PLANE_TOKEN_SALT = "notebooks.sql_v2.data_plane"
_DATA_PLANE_TOKEN_MAX_AGE_SECONDS = 3600

# Rows in the display page the run envelope carries to the UI.
DISPLAY_PAGE_LIMIT = 50
# Rows the kernel fetches and caches per run. Pages within the cache are local slices
# (no Datastore re-query); paging beyond it falls back to a LIMIT/OFFSET re-query.
RESULT_CACHE_ROWS = 300

# The container port the sandbox already exposes (mapped to a host port at create
# time). Mirrors docker_sandbox.AGENT_SERVER_PORT (47821) and
# modal_sandbox.AGENT_SERVER_PORT (8080) — the kernel-server binds it inside the sandbox.
_CONTAINER_PORT_BY_BACKEND: dict[str, int] = {
    KernelRuntime.Backend.DOCKER: 47821,
    KernelRuntime.Backend.MODAL: 8080,
}
_PACKAGE_ROOT = "/tmp/nb_kernel_pkg"
_TARBALL_PATH = "/tmp/nb_kernel.tar.gz"
_SECRET_PATH = "/tmp/nb_sql_v2_secret"
_SERVER_LOG_PATH = "/tmp/nb_kernel_server.log"
_SERVER_PID_PATH = "/tmp/nb_kernel_server.pid"
_SERVER_READY_TIMEOUT_SECONDS = 15
_RUN_POST_TIMEOUT_SECONDS = 10
# A page fetch holds a web worker for the whole kernel -> data plane -> CH round trip.
_PAGE_POST_TIMEOUT_SECONDS = 60
# An interrupt only sets a cancel event and sends a SIGINT in the sandbox: near-instant.
_INTERRUPT_POST_TIMEOUT_SECONDS = 5
# Safety-net TTL for the per-user page-fetch lock: sized just past the POST timeout so a
# worker killed before its finally-release can't wedge the user's paging for long.
PAGE_LOCK_TTL_SECONDS = _PAGE_POST_TIMEOUT_SECONDS + 10
_COMMAND_TOKEN_TTL_SECONDS = 300


def sql_v2_page_lock_key(team_id: int, user_id: int | None) -> str:
    """One in-flight page fetch per user: each holds a web worker for up to the POST timeout."""
    return f"sqlv2_page_inflight:{team_id}:{user_id if user_id is not None else 'anon'}"


class SQLV2KernelNotRunning(Exception):
    """Raised when a run is dispatched but the notebook has no running sandbox."""


class SQLV2PageError(Exception):
    """A page fetch the kernel rejected or failed; message is user-facing."""


def is_sql_v2_enabled(user: User | None) -> bool:
    if user is None or not user.distinct_id:
        return False
    kwargs: dict = {"only_evaluate_locally": False, "send_feature_flag_events": False}
    org = getattr(user, "organization", None)
    if org is not None:
        org_id = str(org.id)
        kwargs["groups"] = {"organization": org_id}
        kwargs["group_properties"] = {"organization": {"id": org_id}}
    return bool(hanzo_insights.feature_enabled(REVAMPED_PY_NOTEBOOKS_FLAG, user.distinct_id, **kwargs))


def mint_callback_token(run_id: str, team_id: int) -> str:
    return signing.dumps({"run_id": run_id, "team_id": team_id}, salt=_CALLBACK_TOKEN_SALT)


def verify_callback_token(token: str) -> tuple[str, int]:
    """Return (run_id, team_id) from a valid token, else raise signing.BadSignature."""
    data = signing.loads(token, salt=_CALLBACK_TOKEN_SALT, max_age=_CALLBACK_TOKEN_MAX_AGE_SECONDS)
    return str(data["run_id"]), int(data["team_id"])


def kernel_server_secret(runtime_id: str) -> str:
    """Per-kernel command-auth secret, derived from Django's signing key.

    Only the backend can derive it (holds SECRET_KEY); a copy is written into the
    sandbox at bootstrap so the kernel-server can verify command tokens. Nothing to
    persist. This is the HMAC analogue of Code's RS256 connection-JWT keypair;
    hardening to RS256 would let the sandbox hold only a public key.
    """
    return hmac.new(
        settings.SECRET_KEY.encode(),
        f"nb-sql-v2-kernel:{runtime_id}".encode(),
        hashlib.sha256,
    ).hexdigest()


def mint_command_token(secret: str, run_id: str, ttl_seconds: int = _COMMAND_TOKEN_TTL_SECONDS) -> str:
    """Sign a short-lived, run-scoped command token; `kernel.auth` verifies it in the sandbox."""
    exp = int(time.time()) + ttl_seconds
    signature = hmac.new(secret.encode(), f"{run_id}.{exp}".encode(), hashlib.sha256).hexdigest()
    return f"{run_id}.{exp}.{signature}"


def _backend_base_url() -> str:
    # The sandbox reaches the host backend here. SANDBOX_API_URL overrides (e.g. ngrok for Modal
    # from local dev). In dev the kernel runs in local Docker, where the host is
    # host.docker.internal (SITE_URL would be an unreachable localhost); in prod the kernel runs
    # remotely and must use the public SITE_URL.
    if settings.SANDBOX_API_URL:
        base = settings.SANDBOX_API_URL
    elif settings.DEBUG:
        base = "http://host.docker.internal:8000"
    else:
        base = settings.SITE_URL
    return base.rstrip("/")


def build_callback_url(run_id: str) -> str:
    return f"{_backend_base_url()}/internal/notebooks/runs/{run_id}/result/"


def build_data_plane_url() -> str:
    return f"{_backend_base_url()}/internal/notebooks/data_plane/query/"


def mint_data_plane_token(notebook_short_id: str, team_id: int, user_id: int | None) -> str:
    return signing.dumps(
        {"notebook_short_id": notebook_short_id, "team_id": team_id, "user_id": user_id},
        salt=_DATA_PLANE_TOKEN_SALT,
    )


def verify_data_plane_token(token: str) -> tuple[str, int, int | None]:
    """Return (notebook_short_id, team_id, user_id) from a valid token, else raise signing.BadSignature."""
    data = signing.loads(token, salt=_DATA_PLANE_TOKEN_SALT, max_age=_DATA_PLANE_TOKEN_MAX_AGE_SECONDS)
    user_id = data.get("user_id")
    return str(data["notebook_short_id"]), int(data["team_id"]), int(user_id) if user_id is not None else None


def _find_running_runtime(notebook: Notebook, user: User | None) -> KernelRuntime | None:
    runtime = (
        KernelRuntime.objects.filter(
            team_id=notebook.team_id,
            notebook_short_id=notebook.short_id,
            user=user if isinstance(user, User) else None,
        )
        .order_by("-last_used_at")
        .first()
    )
    if runtime is None or not runtime.sandbox_id:
        return None
    if runtime.status not in (KernelRuntime.Status.RUNNING, KernelRuntime.Status.STARTING):
        return None
    return runtime


_COMMAND_TOKEN_HEADER = "X-Command-Token"


def _sandbox_auth_headers(connect_token: str | None, command_token: str | None = None) -> dict[str, str]:
    """Auth headers for a request into the sandbox.

    The Modal connect token rides `Authorization: Bearer` — Modal's edge accepts it there
    as well as in the query string, and headers stay out of proxy/access logs while URLs
    do not. With that slot taken, the kernel's command token travels in its own header.
    """
    headers: dict[str, str] = {}
    if connect_token:
        headers["Authorization"] = f"Bearer {connect_token}"
    if command_token:
        headers[_COMMAND_TOKEN_HEADER] = command_token
    return headers


def _server_version(server_url: str, connect_token: str | None) -> str | None:
    """The deployed package hash the running server reports, or None if unreachable."""
    try:
        response = requests.get(
            f"{server_url.rstrip('/')}/health",
            headers=_sandbox_auth_headers(connect_token),
            timeout=2,
        )
        if response.status_code != 200:
            return None
        return str(response.json().get("version") or "")
    except (requests.RequestException, ValueError):
        return None


def _wait_for_server_ready(server_url: str, connect_token: str | None, expected_version: str) -> None:
    deadline = time.monotonic() + _SERVER_READY_TIMEOUT_SECONDS
    while time.monotonic() < deadline:
        if _server_version(server_url, connect_token) == expected_version:
            return
        time.sleep(0.3)
    raise RuntimeError("SQLV2 kernel-server did not become ready")


def _deploy_kernel_server(sandbox: SandboxBase, runtime: KernelRuntime, package: bytes, version: str) -> None:
    """Write the kernel package + secret into the sandbox and (re)launch the server.

    This is the only place the control plane (write_file/execute) is used, exactly
    as Code bootstraps its agent-server. Per-run dispatch is a plain authed POST.
    """
    port = _CONTAINER_PORT_BY_BACKEND.get(runtime.backend, 47821)
    sandbox.write_file(_SECRET_PATH, kernel_server_secret(str(runtime.id)).encode())
    sandbox.write_file(_TARBALL_PATH, package)
    # Stop a previous server via its PID file — never pkill by our own name: the
    # pattern would match this very launch command's shell and kill it mid-deploy.
    # The pkill lines only clear pre-package servers (distinct names, safe to
    # match) from sandboxes that predate the PID file; drop them once those age out.
    # `< /dev/null` detaches the server from the exec's pipes so `execute` returns.
    # The backgrounded launch must stay a single simple command (PYTHONPATH= prefix,
    # no `cd X && …`) — a compound list backgrounds a wrapper subshell, so `$!`
    # would record the wrapper's PID and the next redeploy would kill nothing.
    # Prefer the notebook venv python (has pyarrow).
    launch = (
        f"kill $(cat {_SERVER_PID_PATH} 2>/dev/null) 2>/dev/null || true; "
        "pkill -f '[n]b_sql_v2_kernel_server' 2>/dev/null; pkill -f '[n]b_data_v2_kernel_server' 2>/dev/null; "
        f"rm -rf {_PACKAGE_ROOT} && mkdir -p {_PACKAGE_ROOT} && tar -xzf {_TARBALL_PATH} -C {_PACKAGE_ROOT} && "
        'PY=/opt/notebook-venv/bin/python3; [ -x "$PY" ] || PY=python3; '
        f'PYTHONPATH={_PACKAGE_ROOT} nohup "$PY" -m {SANDBOX_PACKAGE_NAME}.server '
        f"--port {port} --secret-file {_SECRET_PATH} --version {version} "
        f"> {_SERVER_LOG_PATH} 2>&1 < /dev/null & echo $! > {_SERVER_PID_PATH}"
    )
    sandbox.execute(launch, timeout_seconds=30)


def ensure_sql_v2_server(notebook: Notebook, user: User | None) -> KernelRuntime:
    """Ensure the in-sandbox kernel-server is running the current package version.

    Idempotent — a healthy server at the expected version is reused as-is; a stale
    or unreachable one is redeployed from the freshly built tarball (this is the
    dev loop: edit `kernel/`, next run redeploys, no image rebuild).
    """
    runtime = _find_running_runtime(notebook, user)
    if runtime is None:
        raise SQLV2KernelNotRunning()

    package, version = kernel_package_bytes_and_hash()
    if runtime.server_url and _server_version(runtime.server_url, runtime.server_connect_token) == version:
        return runtime

    sandbox_class = get_sandbox_class_for_backend(runtime.backend)
    assert runtime.sandbox_id  # _find_running_runtime only returns runtimes with a sandbox
    sandbox = sandbox_class.get_by_id(runtime.sandbox_id)
    _deploy_kernel_server(sandbox, runtime, package, version)

    credentials = sandbox.get_connect_credentials()
    _wait_for_server_ready(credentials.url, credentials.token, version)

    runtime.server_url = credentials.url
    runtime.server_connect_token = credentials.token
    # A redeploy relaunches the kernel inside a live sandbox and keeps this row, so the frame
    # snapshot would outlive the kernel that produced it — the one path where the row's
    # lifetime stops tracking the kernel's. Its registrations are gone with the namespace, so
    # drop the snapshot; the next run repopulates it from the new kernel's catalog.
    runtime.frames = None
    runtime.save(update_fields=["server_url", "server_connect_token", "frames"])
    return runtime


def dispatch_sql_v2_run(
    notebook: Notebook,
    user: User | None,
    run: NotebookNodeRun,
    code: str,
    node_type: str = "insightsql",
    output_name: str = "",
    inputs: list[dict] | None = None,
) -> None:
    """Dispatch a run to the in-sandbox kernel-server with a single authed HTTP POST.

    Returns as soon as the server accepts (202); the result arrives via the callback. A
    kernel node (python or duckdb) carries the `node`/`inputs` shape the executor consumes;
    a insightsql node keeps the flat `code` the capped-fetch path reads — the paths stay additive.
    """
    runtime = ensure_sql_v2_server(notebook, user)
    assert runtime.server_url  # ensure_sql_v2_server always returns a runtime with a live server_url
    # Record which kernel took the run: the sandbox can't name it (its secret is a one-way
    # derivation of the runtime id, not the id), and letting it name one would let a sandbox
    # write over another kernel's state. So the backend decides here and the callback follows.
    run.kernel_runtime_id = runtime.id
    run.save(update_fields=["kernel_runtime_id", "updated_at"])
    command_token = mint_command_token(kernel_server_secret(str(runtime.id)), str(run.id))
    user_id = user.id if isinstance(user, User) else None
    payload: dict = {
        "run_id": str(run.id),
        "callback_url": build_callback_url(str(run.id)),
        "callback_token": mint_callback_token(str(run.id), notebook.team_id),
        "data_plane_url": build_data_plane_url(),
        "data_plane_token": mint_data_plane_token(notebook.short_id, notebook.team_id, user_id),
        "page_limit": DISPLAY_PAGE_LIMIT,
        "cache_limit": RESULT_CACHE_ROWS,
    }
    if node_type in ("python", "duckdb"):
        payload["node"] = {"type": node_type, "code": code, "output_name": output_name}
        payload["inputs"] = inputs or []
    else:
        payload["code"] = code
    response = requests.post(
        f"{runtime.server_url.rstrip('/')}/run",
        json=payload,
        headers=_sandbox_auth_headers(runtime.server_connect_token, command_token),
        timeout=_RUN_POST_TIMEOUT_SECONDS,
    )
    response.raise_for_status()


def fetch_sql_v2_page(notebook: Notebook, user: User | None, run: NotebookNodeRun, offset: int, limit: int) -> dict:
    """Fetch one result page through the kernel — a bounded synchronous read.

    A insightsql run pages by re-querying its stored code through the data plane with
    LIMIT/OFFSET; a kernel run (python/duckdb) pages by slicing its on-sandbox result
    frame, keyed by `result_id` — its code is not a InsightsQL query, so no data-plane
    fallback exists for it. Unlike a run this never bootstraps the server (no control
    plane from a web worker); a missing or unreachable server means the user has to re-run.
    """
    runtime = _find_running_runtime(notebook, user)
    if runtime is None or not runtime.server_url:
        raise SQLV2KernelNotRunning()

    command_token = mint_command_token(kernel_server_secret(str(runtime.id)), str(run.id))
    user_id = user.id if isinstance(user, User) else None
    payload: dict = {
        "run_id": str(run.id),
        "offset": offset,
        "limit": limit,
    }
    if run.node_type == NotebookNodeRun.NodeType.INSIGHTSQL:
        payload["code"] = run.code
        payload["data_plane_url"] = build_data_plane_url()
        payload["data_plane_token"] = mint_data_plane_token(notebook.short_id, notebook.team_id, user_id)
    else:
        payload["result_id"] = str(run.result_id)
    try:
        response = requests.post(
            f"{runtime.server_url.rstrip('/')}/page",
            json=payload,
            headers=_sandbox_auth_headers(runtime.server_connect_token, command_token),
            timeout=_PAGE_POST_TIMEOUT_SECONDS,
        )
    except requests.RequestException as exc:
        raise SQLV2KernelNotRunning() from exc

    if response.status_code in (404, 401, 403):
        # 404 → kernel-server predates the /page route; 401/403 → the command token expired or
        # the kernel was redeployed with a new secret. All are session-level, not query errors —
        # a re-run redeploys/reissues, so surface them as "re-run" (503), not a page failure.
        raise SQLV2KernelNotRunning()
    if response.status_code == 400:
        raise SQLV2PageError(_kernel_error_detail(response))
    if response.status_code != 200:
        # Any other non-200 (e.g. a kernel 500) is an infrastructure problem, not a bad query.
        raise SQLV2KernelNotRunning()
    return response.json()


def interrupt_sql_v2_run(notebook: Notebook, user: User | None, run: NotebookNodeRun) -> bool:
    """Ask the kernel-server to interrupt a run; return whether the kernel knew the run.

    False means the run never reached the kernel (dispatch still in flight) or already
    finished there; the caller decides how to surface that. Raises SQLV2KernelNotRunning
    when no reachable kernel exists at all, in which case the run's callback can never
    arrive and the caller may mark the run terminal itself.
    """
    runtime = _find_running_runtime(notebook, user)
    if runtime is None or not runtime.server_url:
        raise SQLV2KernelNotRunning()

    command_token = mint_command_token(kernel_server_secret(str(runtime.id)), str(run.id))
    try:
        response = requests.post(
            f"{runtime.server_url.rstrip('/')}/interrupt",
            json={"run_id": str(run.id)},
            headers=_sandbox_auth_headers(runtime.server_connect_token, command_token),
            timeout=_INTERRUPT_POST_TIMEOUT_SECONDS,
        )
    except requests.RequestException as exc:
        raise SQLV2KernelNotRunning() from exc
    if response.status_code != 200:
        raise SQLV2KernelNotRunning()
    try:
        body = response.json()
    except ValueError:
        return True
    # A pre-run-scoped kernel-server omits `known`; treat its interrupt as delivered.
    return bool(body.get("known", True))


def _kernel_error_detail(response: requests.Response) -> str:
    try:
        detail = response.json().get("error")
        if isinstance(detail, str) and detail:
            return detail
    except ValueError:
        pass
    return "Page fetch failed in the sandbox."