import asyncio import csv import io import json from datetime import datetime, timezone import aiosqlite from fastapi import APIRouter, Depends, HTTPException, Query, Request from fastapi.responses import HTMLResponse, StreamingResponse from sse_starlette.sse import EventSourceResponse from app import poller from app.config import settings from app.database import get_db from app.models import ( BurninJobResponse, BurninStageResponse, CancelBurninRequest, DriveResponse, SmartTestState, StartBurninRequest, UnlockPoolDriveRequest, UpdateDriveRequest, ) from app.renderer import templates # Helpers shared with the extracted sub-routers — keep the underscore- # prefixed local names that existing in-file callers reach for. from ._helpers import ( client_ip as _client_ip, is_stale as _is_stale, operator_for as _operator_for, secret_status as _secret_status, stale_context as _stale_context, SECRET_FIELDS as _SECRET_FIELDS, ) router = APIRouter() # Sub-routers extracted as part of the routes/ package split (1.0.0-34). # Their endpoints get registered against the same APIRouter, so the # external `from app.routes import router` import in app/main.py keeps # working unchanged. Future slices can extract more — drives, burnin, # settings, history — using the same pattern. # # Absolute imports (`import app.routes.X as _Y`) instead of relative # (`from . import X as _Y`) so we stay safe even if a future top-level # `from app import X` is reintroduced here — `from app import auth` # would bind `auth` on the `app.routes` package namespace and shadow # any relative-submodule lookup. Absolute imports always resolve to # `app.routes.X` regardless of what's already bound on the package. import app.routes.auth as _auth_routes # noqa: E402 import app.routes.system as _system_routes # noqa: E402 import app.routes.history as _history_routes # noqa: E402 import app.routes.audit as _audit_routes # noqa: E402 import app.routes.stats as _stats_routes # noqa: E402 import app.routes.report as _report_routes # noqa: E402 import app.routes.settings as _settings_routes # noqa: E402 import app.routes.drives as _drives_routes # noqa: E402 import app.routes.burnin as _burnin_routes # noqa: E402 router.include_router(_auth_routes.router) router.include_router(_system_routes.router) router.include_router(_history_routes.router) router.include_router(_audit_routes.router) router.include_router(_stats_routes.router) router.include_router(_report_routes.router) router.include_router(_settings_routes.router) router.include_router(_drives_routes.router) router.include_router(_burnin_routes.router) # Drives helpers — re-exported for the dashboard + SSE handlers in this # file AND for `from app.routes import _fetch_drives_for_template` # from mailer.py (existing back-compat shim). from ._drives_helpers import ( # noqa: E402 _DRIVES_QUERY, _row_to_drive, _build_smart, _compute_status, _compute_eta_seconds, _eta_seconds, _fetch_burnin_by_drive, _fetch_drives_for_template, ) # _stale_context is now imported from ._helpers above. # --------------------------------------------------------------------------- # Dashboard # --------------------------------------------------------------------------- @router.get("/", response_class=HTMLResponse) async def dashboard(request: Request, db: aiosqlite.Connection = Depends(get_db)): drives = await _fetch_drives_for_template(db) ps = poller.get_state() return templates.TemplateResponse(request, "dashboard.html", { "request": request, "drives": drives, "poller": ps, **_stale_context(ps), }) # --------------------------------------------------------------------------- # SSE — live drive table updates # --------------------------------------------------------------------------- @router.get("/sse/drives") async def sse_drives(request: Request): q = poller.subscribe() async def generate(): try: while True: # Wait for next poll notification or keepalive timeout try: payload = await asyncio.wait_for(q.get(), timeout=25.0) except asyncio.TimeoutError: if await request.is_disconnected(): break yield {"event": "keepalive", "data": ""} continue if await request.is_disconnected(): break # Extract alert from payload (may be None for regular polls) alert = None if isinstance(payload, dict): alert = payload.get("alert") # Render fresh table HTML async with aiosqlite.connect(settings.db_path) as db: db.row_factory = aiosqlite.Row await db.execute("PRAGMA journal_mode=WAL") drives = await _fetch_drives_for_template(db) html = templates.env.get_template( "components/drives_table.html" ).render(drives=drives) yield {"event": "drives-update", "data": html} # Push system sensor state so JS can update temp chips live ps = poller.get_state() yield { "event": "system-sensors", "data": json.dumps({ "system_temps": ps.get("system_temps", {}), "thermal_pressure": ps.get("thermal_pressure", "ok"), "temp_warn_c": settings.temp_warn_c, "temp_crit_c": settings.temp_crit_c, }), } # Push browser notification event if this was a job completion if alert: yield {"event": "job-alert", "data": json.dumps(alert)} finally: poller.unsubscribe(q) return EventSourceResponse(generate()) # --------------------------------------------------------------------------- # JSON API # ---------------------------------------------------------------------------