refactor composer to use sqlalchemy, avoid string concatenations and follow conventions
This commit is contained in:
@@ -15,24 +15,59 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Any
|
||||
import logging
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, IO
|
||||
|
||||
from api.analyses.types import Finding
|
||||
from api.runtime.context import current_analysis, current_run_id
|
||||
|
||||
logger = logging.getLogger("nvi.runtime.events")
|
||||
|
||||
# run_id -> queue of events. Events are plain dicts.
|
||||
_queues: dict[str, asyncio.Queue[dict[str, Any] | None]] = {}
|
||||
|
||||
# run_id -> append-mode JSONL log file. Mirrors every published event so the
|
||||
# whole run is replayable from disk. Path: `<LOG_DIR>/<run_id>.jsonl`.
|
||||
_log_files: dict[str, IO[str]] = {}
|
||||
LOG_DIR = Path(".data/logs")
|
||||
|
||||
|
||||
def log_path(run_id: str) -> Path:
|
||||
"""Where this run's JSONL log lives. Resolved relative to cwd of the api
|
||||
process — `/app/.data/logs/` in the pod."""
|
||||
return LOG_DIR / f"{run_id}.jsonl"
|
||||
|
||||
|
||||
# ── Queue management ──
|
||||
|
||||
def open_run(run_id: str) -> asyncio.Queue:
|
||||
q: asyncio.Queue = asyncio.Queue()
|
||||
_queues[run_id] = q
|
||||
try:
|
||||
LOG_DIR.mkdir(parents=True, exist_ok=True)
|
||||
_log_files[run_id] = log_path(run_id).open("a", encoding="utf-8")
|
||||
except OSError as e:
|
||||
# Filesystem issues shouldn't kill the run; just lose the on-disk log.
|
||||
logger.warning("could not open log file for run %s: %s", run_id, e)
|
||||
return q
|
||||
|
||||
|
||||
def _write_log(run_id: str, event: dict[str, Any]) -> None:
|
||||
f = _log_files.get(run_id)
|
||||
if f is None:
|
||||
return
|
||||
try:
|
||||
record = {"t": time.time(), **event}
|
||||
f.write(json.dumps(record, default=str) + "\n")
|
||||
f.flush() # so `kubectl exec ... -- tail -f` shows live progress
|
||||
except (OSError, ValueError) as e:
|
||||
logger.warning("log write failed for run %s: %s", run_id, e)
|
||||
|
||||
|
||||
async def publish(run_id: str, event: dict[str, Any]) -> None:
|
||||
_write_log(run_id, event)
|
||||
q = _queues.get(run_id)
|
||||
if q is not None:
|
||||
await q.put(event)
|
||||
@@ -54,6 +89,12 @@ async def close(run_id: str) -> None:
|
||||
|
||||
def drop(run_id: str) -> None:
|
||||
_queues.pop(run_id, None)
|
||||
f = _log_files.pop(run_id, None)
|
||||
if f is not None:
|
||||
try:
|
||||
f.close()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
# ── SSE formatting + streaming ──
|
||||
|
||||
Reference in New Issue
Block a user