stack/src/bcda/ui/runner.py

"""Background runs for the BCDA manager UI.

Streamlit re-executes the script on every interaction, so anything that
takes longer than a click (a sandbox export, a 2 GB download, a flatten
+ load) runs in a daemon thread and is tracked here, in a process-level
registry that survives reruns.  Each run captures the ``bcda`` logger
lines emitted on its own thread (job started, progress, files stored),
so the page can show what the client is doing without the client
knowing about the UI.

    run = runner.submit("export Patient", store.export, client, types=["Patient"])
    ...
    runner.get(run.id).state   # queued · running · done · failed
    runner.get(run.id).log     # ["Export ABC started", "Job in progress: 35%", ...]
"""

from __future__ import annotations

import logging
import threading
import time
import traceback
import uuid
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any

_LOG_LIMIT = 400


@dataclass
class Run:
    """One background job and what it has said so far."""

    id: str
    name: str
    kind: str = "other"
    state: str = "queued"
    started: float = field(default_factory=time.time)
    finished: float | None = None
    result: Any = None
    error: str | None = None
    log: list[str] = field(default_factory=list)
    thread_id: int | None = None

    @property
    def elapsed(self) -> float:
        return (self.finished or time.time()) - self.started

    @property
    def active(self) -> bool:
        return self.state in ("queued", "running")

    def append(self, line: str) -> None:
        self.log.append(line)
        if len(self.log) > _LOG_LIMIT:
            del self.log[: len(self.log) - _LOG_LIMIT]


class _ThreadLogHandler(logging.Handler):
    """Route ``bcda.*`` log records to the Run owning the emitting thread."""

    def __init__(self, registry: Runner) -> None:
        super().__init__(level=logging.DEBUG)
        self._registry = registry

    def emit(self, record: logging.LogRecord) -> None:
        run = self._registry._by_thread.get(record.thread)
        if run is not None:
            run.append(record.getMessage())


class Runner:
    """Process-level registry of background runs."""

    def __init__(self) -> None:
        self._runs: dict[str, Run] = {}
        self._by_thread: dict[int, Run] = {}
        self._lock = threading.Lock()
        self._handler: _ThreadLogHandler | None = None

    # ── logging capture ──────────────────────────────────────

    def _ensure_handler(self) -> None:
        if self._handler is None:
            self._handler = _ThreadLogHandler(self)
            logger = logging.getLogger("bcda")
            logger.addHandler(self._handler)
            if logger.level == logging.NOTSET or logger.level > logging.DEBUG:
                logger.setLevel(logging.DEBUG)

    # ── API ──────────────────────────────────────────────────

    def submit(
        self,
        name: str,
        fn: Callable[..., Any],
        *args: Any,
        kind: str = "other",
        **kwargs: Any,
    ) -> Run:
        """Start *fn* on a daemon thread; return its Run immediately.

        ``kind`` labels the run for metrics (export, adopt, load,
        vendor_check, other) — see :mod:`bcda.ui.metrics`.
        """
        from bcda.ui import metrics

        self._ensure_handler()
        run = Run(id=uuid.uuid4().hex[:8], name=name, kind=kind)
        with self._lock:
            self._runs[run.id] = run

        def _target() -> None:
            run.thread_id = threading.get_ident()
            with self._lock:
                self._by_thread[run.thread_id] = run
            run.state = "running"
            run.append(f"started: {name}")
            metrics.run_started(kind)
            # The state flips last: ``wait()`` and the page read ``state``
            # without a lock, so result, log and finish time must already be
            # in place when they see done/failed.
            try:
                run.result = fn(*args, **kwargs)
                run.append("finished")
                run.finished = time.time()
                run.state = "done"
            except Exception as exc:  # noqa: BLE001 — surfaced to the page
                run.error = f"{type(exc).__name__}: {exc}"
                run.append(run.error)
                run.append(traceback.format_exc(limit=3))
                run.finished = time.time()
                run.state = "failed"
            finally:
                metrics.run_finished(kind, run.state, run.elapsed)
                with self._lock:
                    self._by_thread.pop(run.thread_id, None)

        threading.Thread(target=_target, name=f"bcda-ui-{run.id}", daemon=True).start()
        return run

    def get(self, run_id: str) -> Run | None:
        return self._runs.get(run_id)

    def all(self) -> list[Run]:
        """Newest first."""
        return sorted(self._runs.values(), key=lambda r: r.started, reverse=True)

    def active(self) -> list[Run]:
        return [r for r in self.all() if r.active]

    def wait(self, run_id: str, timeout: float = 30.0) -> Run:
        """Block until the run leaves queued/running (tests and the CLI)."""
        run = self._runs[run_id]
        deadline = time.time() + timeout
        while run.active and time.time() < deadline:
            time.sleep(0.02)
        return run

    def clear_finished(self) -> int:
        with self._lock:
            gone = [k for k, r in self._runs.items() if not r.active]
            for k in gone:
                del self._runs[k]
        return len(gone)


runner = Runner()
"""The registry the UI and CLI share."""