# #region AppModule [C:5] [TYPE Module] [SEMANTICS fastapi, websocket, startup, scheduler, middleware] # @BRIEF The main entry point for the FastAPI application. # @LAYER API # @RELATION DEPENDS_ON -> [ApiRoutesModule] # @INVARIANT All WebSocket connections must be properly cleaned up on disconnect. # @INVARIANT All WebSocket connections must be authenticated via JWT or API key token (see [SEC:C-3]). # @PRE Python environment and dependencies installed; configuration database available. # @POST FastAPI app instance is created, middleware configured, and routes registered. # @SIDE_EFFECT Starts background scheduler and binds network ports for HTTP/WS traffic. # @DATA_CONTRACT [HTTP Request | WS Message] -> [HTTP Response | JSON Log Stream] # @RATIONALE Duplicate @RELATION and @INVARIANT lines removed from header. Import sorting unified via ruff isort (I) rule across src/ — 139 fixes applied. import asyncio from contextlib import asynccontextmanager import os from pathlib import Path # project_root is used for static files mounting project_root = Path(__file__).resolve().parent.parent.parent from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect, status from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import FileResponse, JSONResponse from fastapi.staticfiles import StaticFiles from starlette.middleware.base import BaseHTTPMiddleware from starlette.middleware.sessions import SessionMiddleware from typing import Any from .api import auth from .api.routes import ( admin, admin_api_keys, assistant, clean_release, clean_release_v2, dashboards, dataset_review, datasets, environments, git, health, llm, maintenance, mappings, migration, plugins, profile, reports, settings, storage, tasks, translate, ) from .api.routes.validation_tasks import router as validation_tasks from .core.auth.security import get_password_hash from .core.cot_logger import seed_trace_id from .core.database import AuthSessionLocal, init_db from .core.encryption_key import ensure_encryption_key from .core.logger import belief_scope, logger from .core.utils.network import NetworkError from .dependencies import get_scheduler_service, get_task_manager from .models.auth import Role, User # #region lifespan [C:3] [TYPE Function] # @BRIEF Async context manager for FastAPI startup/shutdown lifecycle. # @RELATION CALLS -> [init_db] # @RELATION CALLS -> [AppDependencies] # @POST On startup: admin exists, scheduler started. On shutdown: scheduler stopped. # @RATIONALE Alembic migrations removed from lifespan — they now run exclusively # in docker/backend.entrypoint.sh (wait_for_db → alembic upgrade head). # Running migrations in both places added ~5s startup overhead and masked # partial failures. init_db() remains as a safety net for tables without # dedicated Alembic migrations (e.g., newly added models during development). @asynccontextmanager async def lifespan(app: FastAPI): # Startup seed_trace_id() with belief_scope("startup_event"): logger.info("🔐 Ensuring encryption key...") ensure_encryption_key() logger.info("🗄️ Initializing database tables (safety net)...") init_db() logger.info("👤 Bootstrapping admin user...") ensure_initial_admin_user() # Clean up stuck validation runs from previous backend lifetime # (runs left "running" in DB when the in-memory task queue was lost) try: from sqlalchemy.orm import Session as _Ses from src.core.database import SessionLocal as _Db from src.models.llm import ValidationRun as _VR from datetime import datetime, timezone _s: _Ses = _Db() _stuck = _s.query(_VR).filter(_VR.status == "running").all() for _r in _stuck: _r.status = "FAIL" _r.finished_at = datetime.now(timezone.utc) _r.summary = "Force-stopped: backend restarted while run was in progress" logger.reason( f"Force-stopped stuck run {_r.id} (policy={_r.policy_id})", extra={"src": "app.startup", "run_id": _r.id}, ) _s.commit() _s.close() except Exception as _e: logger.warning(f"Failed to clean up stuck validation runs: {_e}") logger.info("⏰ Starting scheduler...") scheduler = get_scheduler_service() scheduler.start() logger.info("✅ Application startup complete") yield # Shutdown scheduler.stop() # #endregion lifespan # #region FastAPI_App [C:3] [TYPE Global] [SEMANTICS app, fastapi, instance, route-registry] # @BRIEF Canonical FastAPI application instance for route, middleware, and websocket registration. # @RELATION DEPENDS_ON -> [ApiRoutesModule] # @RELATION BINDS_TO -> [API_Routes] app = FastAPI( title="Superset Tools API", description="API for managing Superset automation tools and plugins.", version="1.0.0", lifespan=lifespan, ) # TraceContextMiddleware is a raw ASGI middleware (not BaseHTTPMiddleware). # This is by design — BaseHTTPMiddleware (Starlette 0.50.0) uses # anyio.create_task_group() internally, creating separate asyncio tasks for # dispatch vs call_next. ContextVars set in dispatch() are NOT visible to # outer middleware layers (like log_requests). Raw ASGI middleware runs in # the root task context, making trace_id visible to ALL middleware layers. # See trace.py @RATIONALE for details. from .core.middleware.trace import TraceContextMiddleware # noqa: E402 app.add_middleware(TraceContextMiddleware) # #endregion FastAPI_App # #region ensure_initial_admin_user [C:3] [TYPE Function] # @BRIEF Ensures initial admin user exists when bootstrap env flags are enabled. # @RELATION DEPENDS_ON -> [AuthRepository] def ensure_initial_admin_user() -> None: raw_flag = os.getenv("INITIAL_ADMIN_CREATE", "false").strip().lower() if raw_flag not in {"1", "true", "yes", "on"}: return username = os.getenv("INITIAL_ADMIN_USERNAME", "").strip() password = os.getenv("INITIAL_ADMIN_PASSWORD", "").strip() if not username or not password: logger.warning( "INITIAL_ADMIN_CREATE is enabled but INITIAL_ADMIN_USERNAME/INITIAL_ADMIN_PASSWORD is missing; skipping bootstrap." ) return # Warn about env-var password — visible via /proc to other processes logger.warning( "INITIAL_ADMIN_PASSWORD is set via environment variable. " "This is visible to other processes on the same host. " "Consider removing the env var after bootstrap." ) db = AuthSessionLocal() try: admin_role = db.query(Role).filter(Role.name == "Admin").first() if not admin_role: admin_role = Role(name="Admin", description="System Administrator", is_admin=True) db.add(admin_role) db.commit() db.refresh(admin_role) existing_user = db.query(User).filter(User.username == username).first() if existing_user: logger.info( "Initial admin bootstrap skipped: user '%s' already exists.", username ) return new_user = User( username=username, email=None, password_hash=get_password_hash(password), auth_source="LOCAL", is_active=True, ) new_user.roles.append(admin_role) db.add(new_user) db.commit() logger.info( "Initial admin user '%s' created from environment bootstrap.", username ) except Exception as exc: db.rollback() logger.error("Failed to bootstrap initial admin user: %s", exc) raise finally: db.close() # #endregion ensure_initial_admin_user # #region run_alembic_migrations [C:2] [TYPE Function] # @BRIEF Applies all pending Alembic migrations against DATABASE_URL. # DEPRECATED: Migrations now run exclusively in docker/backend.entrypoint.sh. # Kept for local development (manual invocation). # @POST All Alembic migrations up to 'head' are applied. # @SIDE_EFFECT Executes ALTER TABLE / CREATE TABLE via Alembic chain. # @DEPRECATED Migrations moved to entrypoint.sh — this function is kept for # local development only. Do NOT call from lifespan or production code. def run_alembic_migrations() -> None: from alembic.config import Config as AlembicConfig from alembic import command as alembic_command with belief_scope("run_alembic_migrations"): try: alembic_cfg = AlembicConfig("alembic.ini") alembic_cfg.set_main_option("script_location", "alembic") alembic_command.upgrade(alembic_cfg, "head") logger.reason("Alembic migrations applied up to head") except Exception as exc: logger.explore( "Alembic migration failed — check migration chain or alembic.ini", extra={"error": str(exc)}, ) raise # #endregion run_alembic_migrations # #region app_middleware [TYPE Block] # @BRIEF Configure application-wide middleware (Session, CORS). # @RATIONALE SessionMiddleware uses SESSION_SECRET_KEY independent of JWT SECRET_KEY (see [SEC:H-4]). # CORS allow_origins crashes early if ALLOWED_ORIGINS is unset — no "*" fallback (see [SEC:M-1]). # @REJECTED Hardcoded allow_origins=["*"] rejected — open CORS allows any origin to access the API, # which is a Class 1 security violation in production. # @REJECTED SessionMiddleware sharing JWT SECRET_KEY rejected in [SEC:H-4] — key reuse expands blast radius. # Configure Session Middleware (required by Authlib for OAuth2 flow) from .core.auth.config import auth_config _session_secret = os.getenv("SESSION_SECRET_KEY", "").strip() if not _session_secret: _session_secret = auth_config.SECRET_KEY logger.warning( "SESSION_SECRET_KEY not set. Falling back to AUTH_SECRET_KEY. " "Set a separate SESSION_SECRET_KEY in .env for production isolation." ) app.add_middleware(SessionMiddleware, secret_key=_session_secret) # Configure CORS _allowed_origins_raw = os.getenv("ALLOWED_ORIGINS", "").strip() if not _allowed_origins_raw: logger.warning( "ALLOWED_ORIGINS not set. CORS will reject all cross-origin requests. " "Set ALLOWED_ORIGINS to a comma-separated list of allowed origins." ) _allowed_origins = [] else: _allowed_origins = [o.strip() for o in _allowed_origins_raw.split(",") if o.strip()] app.add_middleware( CORSMiddleware, allow_origins=_allowed_origins, allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # HSTS — Strict-Transport-Security header (see [SEC:L-2]) # Only active when FORCE_HTTPS=true (enterprise deployments with proper certs). # In dev or behind HTTP-only proxies, HSTS is disabled to avoid lockout. class HSTSMiddleware(BaseHTTPMiddleware): """Add Strict-Transport-Security header when FORCE_HTTPS=true. Enterprise note: In production, set HSTS at the nginx/ingress level (add_header Strict-Transport-Security ...). This middleware is a fallback for environments where nginx is not configured to do so. """ def __init__(self, app): super().__init__(app) self._enabled = os.getenv("FORCE_HTTPS", "").strip().lower() in {"1", "true", "yes"} async def dispatch(self, request: Request, call_next): response = await call_next(request) if self._enabled: response.headers["Strict-Transport-Security"] = "max-age=31536000; includeSubDomains" return response app.add_middleware(HSTSMiddleware) # #endregion app_middleware # #region global_exception_handler [C:2] [TYPE Function] # @BRIEF Global exception handler — logs all unhandled 500 errors into the app logger. # @PRE request is a FastAPI Request object. # @POST Logs full traceback to superset_tools_app logger; returns 500. # @RATIONALE FastAPI/Starlette writes unhandled exceptions to uvicorn.error logger, # which is not captured in docker logs. This handler ensures every 500 # appears in our structured JSON log output. @app.exception_handler(Exception) async def global_exception_handler(request: Request, exc: Exception): client_host = request.client.host if request.client else "unknown" logger.exception( f"Unhandled exception: {request.method} {request.url.path}", extra={ "src": "app.exception_handler", "path": request.url.path, "method": request.method, "query_params": dict(request.query_params), "client": client_host, }, ) return JSONResponse( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, content={"detail": "Internal server error", "path": request.url.path}, ) # #endregion global_exception_handler # #region network_error_handler [C:1] [TYPE Function] # @BRIEF Global exception handler for NetworkError. # @PRE request is a FastAPI Request object. # @POST Returns 503 HTTP Exception. @app.exception_handler(NetworkError) async def network_error_handler(request: Request, exc: NetworkError): with belief_scope("network_error_handler"): logger.error(f"Network error: {exc}") return HTTPException( status_code=503, detail="Environment unavailable. Please check if the Superset instance is running.", ) # #endregion network_error_handler # #region log_requests [C:3] [TYPE Function] # @BRIEF Middleware to log incoming HTTP requests and their response status. # @RELATION DEPENDS_ON -> [LoggerModule] # @PRE request is a FastAPI Request object. # @POST Logs request and response details. @app.middleware("http") async def log_requests(request: Request, call_next): with belief_scope("log_requests"): # Avoid spamming logs for polling endpoints is_polling = request.url.path.endswith("/api/tasks") and request.method == "GET" if not is_polling: import json as _json, logging as _lg if logger.isEnabledFor(_lg.INFO): logger.info(f"Incoming request: {request.method} {request.url.path}") else: _json.dump({"ts": __import__('datetime').datetime.now(__import__('datetime').UTC).isoformat(), "level": "INFO", "src": "log_requests", "marker": "REASON", "intent": f"Incoming request: {request.method} {request.url.path}"}, sys.stderr, ensure_ascii=False) sys.stderr.write("\n") sys.stderr.flush() try: response = await call_next(request) if not is_polling: if logger.isEnabledFor(_lg.INFO): logger.info(f"Response status: {response.status_code} for {request.url.path}") else: _json.dump({"ts": __import__('datetime').datetime.now(__import__('datetime').UTC).isoformat(), "level": "INFO", "src": "log_requests", "marker": "REFLECT", "intent": f"Response: {response.status_code} for {request.url.path}"}, sys.stderr, ensure_ascii=False) sys.stderr.write("\n") sys.stderr.flush() return response except NetworkError as e: logger.error(f"Network error caught in middleware: {e}") raise HTTPException( status_code=503, detail="Environment unavailable. Please check if the Superset instance is running.", ) # #endregion log_requests # #region API_Routes [C:3] [TYPE Block] # @BRIEF Register all FastAPI route groups exposed by the application entrypoint. # @RELATION DEPENDS_ON -> [FastAPI_App] # @RELATION DEPENDS_ON -> [Route_Group_Contracts] # @RELATION DEPENDS_ON -> [AuthApi] # @RELATION DEPENDS_ON -> [AdminApi] # @RELATION DEPENDS_ON -> [PluginsRouter] # @RELATION DEPENDS_ON -> [TasksRouter] # @RELATION DEPENDS_ON -> [SettingsRouter] # @RELATION DEPENDS_ON -> [ReportsRouter] # @RELATION DEPENDS_ON -> [LlmRoutes] # @RELATION DEPENDS_ON -> [CleanReleaseV2Api] # @RELATION DEPENDS_ON -> [MaintenanceRouter] # Include API routes app.include_router(auth.router) app.include_router(admin.router) app.include_router(admin_api_keys.router) app.include_router(plugins.router, prefix="/api/plugins", tags=["Plugins"]) app.include_router(tasks.router, prefix="/api/tasks", tags=["Tasks"]) app.include_router(settings.router, prefix="/api/settings", tags=["Settings"]) app.include_router(environments.router, tags=["Environments"]) app.include_router(mappings.router, prefix="/api/mappings", tags=["Mappings"]) app.include_router(migration.router) app.include_router(git.router, prefix="/api/git", tags=["Git"]) app.include_router(llm.router, prefix="/api/llm", tags=["LLM"]) app.include_router(storage.router, prefix="/api/storage", tags=["Storage"]) app.include_router(dashboards.router) app.include_router(datasets.router) app.include_router(reports.router) app.include_router(assistant.router, prefix="/api/assistant", tags=["Assistant"]) app.include_router(clean_release.router) app.include_router(clean_release_v2.router) app.include_router(profile.router) app.include_router(dataset_review.router) app.include_router(health.router) app.include_router(translate.router) app.include_router(validation_tasks, prefix="/api/validation-tasks", tags=["Validation Tasks"]) app.include_router(maintenance.maintenance_router) # #endregion API_Routes # #region api.include_routers [C:1] [TYPE Action] [SEMANTICS routes, registration, api] # @BRIEF Registers all API routers with the FastAPI application. # @LAYER API # #endregion api.include_routers # #region _authenticate_websocket [TYPE Function] # @BRIEF Authenticate a WebSocket connection via JWT or API key from query param `token`. # @PRE websocket is a live Starlette WebSocket before accept(). # @POST Returns True if token is valid, logs reason; returns False if rejected. # @RELATION DEPENDS_ON -> [AuthJwtModule] # @RELATION DEPENDS_ON -> [APIKeyModel] # @SIDE_EFFECT Performs DB read to validate API key hash. async def _authenticate_websocket(websocket: WebSocket, endpoint_name: str) -> bool: ws_token = websocket.query_params.get("token", "") if not ws_token: logger.warning( "WebSocket connection rejected: missing token", extra={"endpoint": endpoint_name}, ) return False # Try JWT first, fallback to API key try: from .core.auth.api_key import hash_api_key from .core.auth.jwt import decode_token from .core.database import SessionLocal # Attempt JWT validation payload = decode_token(ws_token) user = payload.get("sub") if isinstance(user, str) and user: logger.reason( "WebSocket authenticated via JWT", extra={"endpoint": endpoint_name, "user": user}, ) return True except Exception: logger.explore("JWT validation failed in WebSocket auth", error="see traceback") pass # Fallback: try API key try: from .core.auth.api_key import hash_api_key from .core.database import SessionLocal from .models.api_key import APIKey key_hash = hash_api_key(ws_token) db = SessionLocal() try: api_key = db.query(APIKey).filter(APIKey.key_hash == key_hash).first() if api_key and api_key.active: logger.reason( "WebSocket authenticated via API key", extra={"endpoint": endpoint_name, "key_name": api_key.name}, ) return True finally: db.close() except Exception: logger.explore("API key validation failed in WebSocket auth", error="see traceback") pass logger.warning( "WebSocket connection rejected: invalid token", extra={"endpoint": endpoint_name}, ) return False # #endregion _authenticate_websocket # #region websocket_endpoint [C:5] [TYPE Function] # @BRIEF Provides a WebSocket endpoint for real-time log streaming of a task with server-side filtering. # @RELATION CALLS -> [TaskManagerPackage] # @RELATION DEPENDS_ON -> [LoggerModule] # @RELATION CALLS -> [_authenticate_websocket] # @PRE task_id must be a valid task ID. WebSocket must be authenticated via `token` query param. # @POST WebSocket connection is managed and logs are streamed until disconnect. # @SIDE_EFFECT Subscribes to TaskManager log queue and broadcasts messages over network. # @DATA_CONTRACT [task_id: str, source: str, level: str] -> [JSON log entry objects] # @INVARIANT Every accepted WebSocket subscription is unsubscribed exactly once even when streaming fails or the client disconnects. # @UX_STATE Connecting -> Streaming -> (Disconnected) # # @TEST_CONTRACT WebSocketLogStreamApi -> # { # required_fields: {websocket: WebSocket, task_id: str}, # optional_fields: {source: str, level: str}, # invariants: [ # "Accepts the WebSocket connection", # "Applies source and level filters correctly to streamed logs", # "Cleans up subscriptions on disconnect" # ] # } # @TEST_FIXTURE valid_ws_connection -> {"task_id": "test_1", "source": "plugin"} # @TEST_EDGE task_not_found_ws -> closes connection or sends error # @TEST_EDGE empty_task_logs -> waits for new logs # @TEST_INVARIANT consistent_streaming -> verifies: [valid_ws_connection] # @TEST_EDGE ws_auth_missing_token -> connection rejected with 4001 # @TEST_EDGE ws_auth_invalid_token -> connection rejected with 4001 @app.websocket("/ws/logs/{task_id}") async def websocket_endpoint( websocket: WebSocket, task_id: str, source: str = None, level: str = None ): """ WebSocket endpoint for real-time log streaming AND task status updates. Sends two message types: - Log entries: plain dicts with level/message/timestamp (backward compatible, no type field) - Status updates: dict with type="task_status" and nested task dict Query Parameters: source: Filter logs by source component (e.g., "plugin", "superset_api") level: Filter logs by minimum level (DEBUG, INFO, WARNING, ERROR) token: JWT or API key for authentication (required, see [SEC:C-3]) """ seed_trace_id() with belief_scope("websocket_endpoint", f"task_id={task_id}"): # ── WebSocket authentication (see [SEC:C-3]) ── if not await _authenticate_websocket(websocket, "ws/logs"): await websocket.close(code=4001, reason="Authentication required") return await websocket.accept() source_filter = source.lower() if source else None level_filter = level.upper() if level else None level_hierarchy = {"DEBUG": 0, "INFO": 1, "WARNING": 2, "ERROR": 3} min_level = level_hierarchy.get(level_filter, 0) if level_filter else 0 logger.reason( "Accepted WebSocket log+status stream connection", extra={ "task_id": task_id, "source_filter": source_filter, "level_filter": level_filter, "min_level": min_level, }, ) task_manager = get_task_manager() log_queue = await task_manager.subscribe_logs(task_id) status_queue = await task_manager.subscribe_status(task_id) logger.reason( "Subscribed WebSocket client to task log and status queues", extra={"task_id": task_id}, ) def matches_filters(log_entry) -> bool: """Check if log entry matches the filter criteria.""" log_source = getattr(log_entry, "source", None) if source_filter and str(log_source or "").lower() != source_filter: return False if level_filter: log_level = level_hierarchy.get(str(log_entry.level).upper(), 0) if log_level < min_level: return False return True async def send_status_update(task: Any) -> None: """Send a structured task status update over the WebSocket.""" status_dict = { "id": task.id, "plugin_id": task.plugin_id, "status": task.status.value if hasattr(task.status, "value") else str(task.status), "started_at": task.started_at.isoformat() if task.started_at else None, "finished_at": task.finished_at.isoformat() if task.finished_at else None, "user_id": task.user_id, "result": task.result, "input_required": task.input_required, "input_request": task.input_request, } await websocket.send_json({ "type": "task_status", "task_id": task.id, "task": status_dict, }) try: # ── Send initial task status ── task = task_manager.get_task(task_id) if task: await send_status_update(task) logger.reason( "Sent initial task status", extra={"task_id": task_id, "status": str(task.status)}, ) # ── Replay initial logs ── logger.reason( "Starting task log stream replay and live forwarding", extra={"task_id": task_id}, ) initial_logs = task_manager.get_task_logs(task_id) initial_sent = 0 for log_entry in initial_logs: if matches_filters(log_entry): log_dict = log_entry.model_dump() log_dict["timestamp"] = log_dict["timestamp"].isoformat() await websocket.send_json(log_dict) initial_sent += 1 logger.reflect( "Initial task log replay completed", extra={ "task_id": task_id, "replayed_logs": initial_sent, "total_available_logs": len(initial_logs), }, ) # ── Send synthetic AWAITING_INPUT prompt if needed ── task = task_manager.get_task(task_id) if task and task.status == "AWAITING_INPUT" and task.input_request: synthetic_log = { "timestamp": task.logs[-1].timestamp.isoformat() if task.logs else "2024-01-01T00:00:00", "level": "INFO", "message": "Task paused for user input (Connection Re-established)", "context": {"input_request": task.input_request}, } await websocket.send_json(synthetic_log) logger.reason( "Replayed awaiting-input prompt to restored WebSocket client", extra={"task_id": task_id, "task_status": task.status}, ) # ── Main loop: listen on both log and status queues ── while True: log_task = asyncio.create_task(log_queue.get()) status_task = asyncio.create_task(status_queue.get()) done, _ = await asyncio.wait( [log_task, status_task], return_when=asyncio.FIRST_COMPLETED, ) for coro in done: result = coro.result() # ── Status update ── if isinstance(result, dict) and result.get("type") == "task_status": await websocket.send_json(result) task_status = result.get("task", {}).get("status", "") if task_status in ("SUCCESS", "FAILED"): logger.reason( "Task reached terminal state via status broadcast; closing stream", extra={"task_id": task_id, "status": task_status}, ) await asyncio.sleep(2) raise StopIteration # exit the while loop continue # ── Log entry ── if not matches_filters(result): continue log_dict = result.model_dump() log_dict["timestamp"] = log_dict["timestamp"].isoformat() await websocket.send_json(log_dict) logger.reflect( "Forwarded task log entry to WebSocket client", extra={ "task_id": task_id, "level": log_dict.get("level"), }, ) if ( "Task completed successfully" in result.message or "Task failed" in result.message ): logger.reason( "Observed terminal task log entry; delaying to preserve client visibility", extra={"task_id": task_id, "message": result.message}, ) await asyncio.sleep(2) except (WebSocketDisconnect, StopIteration) as _ws_exc: if isinstance(_ws_exc, StopIteration): # Task reached terminal state — close cleanly with code 1000 try: await websocket.close(code=1000, reason="Task completed") except Exception: pass logger.reason( "WebSocket client disconnected or stream ended", extra={"task_id": task_id}, ) except Exception as exc: logger.explore( "WebSocket log+status streaming encountered an unexpected failure", extra={"task_id": task_id, "error": str(exc)}, ) raise finally: task_manager.unsubscribe_logs(task_id, log_queue) task_manager.unsubscribe_status(task_id, status_queue) logger.reflect( "Released WebSocket log and status queue subscriptions", extra={"task_id": task_id}, ) # #endregion websocket_endpoint # #region task_events_websocket [C:4] [TYPE Function] # @BRIEF WebSocket endpoint for global task events (status changes for ALL tasks). # @RELATION CALLS -> [TaskManagerPackage] # @PRE WebSocket must be authenticated via `token` query param. # @POST WebSocket streams task status events until disconnect. @app.websocket("/ws/task-events") async def task_events_websocket(websocket: WebSocket): """ WebSocket endpoint for global task events. Streams {type: "task_status", task_id: ..., task: {...}} for ALL task status changes. Query Parameters: token: JWT or API key for authentication (required) """ seed_trace_id() with belief_scope("task_events_websocket"): if not await _authenticate_websocket(websocket, "ws/task-events"): await websocket.close(code=4001, reason="Authentication required") return await websocket.accept() logger.reason("Accepted global task events WebSocket connection") task_manager = get_task_manager() event_queue = await task_manager.subscribe_task_events() logger.reason("Subscribed to global task events") try: while True: event = await event_queue.get() await websocket.send_json(event) logger.reflect( "Forwarded task event to global client", extra={"task_id": event.get("task_id")}, ) except WebSocketDisconnect: logger.reason("Global task events client disconnected") except Exception as exc: logger.explore( "Global task events streaming failed", extra={"error": str(exc)}, ) raise finally: task_manager.unsubscribe_task_events(event_queue) logger.reflect("Released global task events subscription") # #endregion task_events_websocket # #region maintenance_events_websocket [C:4] [TYPE Function] # @BRIEF WebSocket endpoint for maintenance events (created/ended/banner changes). # @RELATION CALLS -> [TaskManagerPackage] # @PRE WebSocket must be authenticated via `token` query param. # @POST WebSocket streams maintenance events until disconnect. @app.websocket("/ws/maintenance/events") async def maintenance_events_websocket(websocket: WebSocket): """ WebSocket endpoint for maintenance events. Streams {type: "maintenance.event_created", ...} / {type: "maintenance.event_ended"}. Query Parameters: token: JWT or API key for authentication (required) """ seed_trace_id() with belief_scope("maintenance_events_websocket"): if not await _authenticate_websocket(websocket, "ws/maintenance/events"): await websocket.close(code=4001, reason="Authentication required") return await websocket.accept() logger.reason("Accepted maintenance events WebSocket connection") task_manager = get_task_manager() event_queue = await task_manager.subscribe_maintenance_events() logger.reason("Subscribed to maintenance events") try: while True: event = await event_queue.get() await websocket.send_json(event) logger.reflect( "Forwarded maintenance event to client", extra={"event_type": event.get("type"), "maintenance_id": event.get("maintenance_id")}, ) except WebSocketDisconnect: logger.reason("Maintenance events client disconnected") except Exception as exc: logger.explore( "Maintenance events streaming failed", extra={"error": str(exc)}, ) raise finally: task_manager.unsubscribe_maintenance_events(event_queue) logger.reflect("Released maintenance events subscription") # #endregion maintenance_events_websocket # #region dataset_websocket_endpoint [C:4] [TYPE Function] # @BRIEF WebSocket endpoint for dataset.updated events — auto-refresh on task completion. # @RELATION CALLS -> [TaskManagerPackage] # @PRE env_id must reference a known environment. # @POST WebSocket streams dataset.updated events until disconnect. # @SIDE_EFFECT Subscribes to dataset event queue in task manager lifecycle. @app.websocket("/ws/datasets/{env_id}") async def dataset_websocket_endpoint(websocket: WebSocket, env_id: str): seed_trace_id() with belief_scope("dataset_websocket_endpoint", f"env_id={env_id}"): # ── WebSocket authentication (see [SEC:C-3]) ── if not await _authenticate_websocket(websocket, "ws/datasets"): await websocket.close(code=4001, reason="Authentication required") return await websocket.accept() logger.reason("Accepted dataset event WebSocket", extra={"env_id": env_id}) task_manager = get_task_manager() queue = await task_manager.subscribe_dataset_events(env_id) try: while True: event = await queue.get() await websocket.send_json(event) logger.reflect("Forwarded dataset.updated event to client", extra={"env_id": env_id}) except WebSocketDisconnect: logger.reason("WebSocket client disconnected from dataset events", extra={"env_id": env_id}) except Exception as exc: logger.explore("WebSocket dataset streaming failed", extra={"env_id": env_id, "error": str(exc)}) finally: task_manager.unsubscribe_dataset_events(env_id, queue) logger.reflect("Released dataset event subscription", extra={"env_id": env_id}) # #endregion dataset_websocket_endpoint # #region translate_run_websocket [C:3] [TYPE Function] # @BRIEF WebSocket endpoint for translation run progress — streams structured status updates. # @PRE run_id must be a valid translation run ID. WebSocket authenticated via `token` query param. # @POST Streams run status JSON every second until terminal state or disconnect. # @SIDE_EFFECT Queries DB each tick for current run status via TranslationOrchestrator. # @UX_STATE Streaming -> Terminal (completed/failed/cancelled) -> Close # @UX_FEEDBACK Client receives {status, total_records, successful_records, failed_records, progressPct, ...} @app.websocket("/ws/translate/run/{run_id}") async def translate_run_websocket(websocket: WebSocket, run_id: str): seed_trace_id() if not await _authenticate_websocket(websocket, "ws/translate/run"): await websocket.close(code=4001, reason="Authentication required") return await websocket.accept() logger.reason("Accepted translate run WebSocket", extra={"run_id": run_id}) try: while True: try: from .core.database import SessionLocal from .plugins.translate.orchestrator_aggregator import TranslationResultAggregator from .plugins.translate.events import TranslationEventLog db = SessionLocal() try: event_log = TranslationEventLog(db) aggregator = TranslationResultAggregator(db, event_log) status = aggregator.get_run_status(run_id) total = status.get("total_records", 0) or 0 done = (status.get("successful_records", 0) or 0) + (status.get("failed_records", 0) or 0) + (status.get("skipped_records", 0) or 0) progress_pct = round((done / total) * 100) if total > 0 else 0 status["progressPct"] = progress_pct await websocket.send_json(status) if status.get("status") in ("COMPLETED", "FAILED", "CANCELLED"): await asyncio.sleep(2) break finally: db.close() except Exception as tick_err: logger.explore("Translate run WS tick error", extra={"run_id": run_id, "error": str(tick_err)}) await websocket.send_json({"error": str(tick_err)}) break await asyncio.sleep(1) except WebSocketDisconnect: logger.reason("Translate run WS disconnected", extra={"run_id": run_id}) except Exception as exc: logger.explore("Translate run WS error", extra={"run_id": run_id, "error": str(exc)}) logger.reflect("Translate run WS closed", extra={"run_id": run_id}) # #endregion translate_run_websocket # #region StaticFiles [C:1] [TYPE Mount] [SEMANTICS static, frontend, spa] # @BRIEF Mounts the frontend build directory to serve static assets. frontend_path = project_root / "frontend" / "build" if frontend_path.exists(): app.mount( "/_app", StaticFiles(directory=str(frontend_path / "_app")), name="static" ) # #region serve_spa [TYPE Function] [C:1] # @PURPOSE Serves the SPA frontend for any path not matched by API routes. # @PRE frontend_path exists. # @POST Returns the requested file or index.html. @app.get("/{file_path:path}", include_in_schema=False) async def serve_spa(file_path: str): with belief_scope("serve_spa"): # Only serve SPA for non-API paths # API routes are registered separately and should be matched by FastAPI first if file_path and ( file_path.startswith("api/") or file_path.startswith("/api/") or file_path == "api" ): # This should not happen if API routers are properly registered # Return 404 instead of serving HTML raise HTTPException( status_code=404, detail=f"API endpoint not found: {file_path}" ) full_path = frontend_path / file_path if file_path and full_path.is_file(): return FileResponse(str(full_path)) return FileResponse(str(frontend_path / "index.html")) # #endregion serve_spa else: # #region read_root [TYPE Function] [C:1] # @PURPOSE A simple root endpoint to confirm that the API is running when frontend is missing. # @PRE None. # @POST Returns a JSON message indicating API status. @app.get("/") async def read_root(): with belief_scope("read_root"): return { "message": "Superset Tools API is running (Frontend build not found)" } # #endregion read_root # #endregion StaticFiles # #endregion AppModule