# #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 os from pathlib import Path # project_root is used for static files mounting project_root = Path(__file__).resolve().parent.parent.parent import asyncio from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import FileResponse from fastapi.staticfiles import StaticFiles from starlette.middleware.base import BaseHTTPMiddleware from starlette.middleware.sessions import SessionMiddleware 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, validation, ) from alembic.config import Config as AlembicConfig from alembic import command as alembic_command 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 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", ) # #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. # @POST All Alembic migrations up to 'head' are applied. # @SIDE_EFFECT Executes ALTER TABLE / CREATE TABLE via Alembic chain. # @RATIONALE Single source of truth for schema changes. All column additions, # drops, renames, and data migrations go through Alembic. # init_db() runs after to create any tables that don't have # a standalone Alembic migration yet. def run_alembic_migrations() -> None: 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 startup_event [C:3] [TYPE Function] # @BRIEF Handles application startup: Alembic → create_all → admin bootstrap → scheduler. # @RELATION CALLS -> [run_alembic_migrations] # @RELATION CALLS -> [init_db] # @POST Schema is up-to-date via Alembic, missing tables created, admin exists, scheduler started. @app.on_event("startup") async def startup_event(): seed_trace_id() with belief_scope("startup_event"): ensure_encryption_key() # 1. Alembic — applies ALL pending migrations (add/drop/rename columns, data migrations) run_alembic_migrations() # 2. init_db — creates new tables via create_all() (does NOT alter existing ones) init_db() ensure_initial_admin_user() scheduler = get_scheduler_service() scheduler.start() # #endregion startup_event # #region shutdown_event [C:3] [TYPE Function] # @BRIEF Handles application shutdown tasks, such as stopping the scheduler. # @RELATION CALLS -> [AppDependencies] # @PRE None. # @POST Scheduler is stopped. # Shutdown event @app.on_event("shutdown") async def shutdown_event(): seed_trace_id() with belief_scope("shutdown_event"): scheduler = get_scheduler_service() scheduler.stop() # #endregion shutdown_event # #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 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: logger.info(f"Incoming request: {request.method} {request.url.path}") try: response = await call_next(request) if not is_polling: logger.info( f"Response status: {response.status_code} for {request.url.path}" ) 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 # Seed trace_id on every request — MUST be the outermost middleware so trace_id # is seeded before log_requests (or any other middleware) runs. from .core.middleware.trace import TraceContextMiddleware app.add_middleware(TraceContextMiddleware) # #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 -> [ValidationRoutes] # @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.router, prefix="/api/validation", tags=["Validation"]) 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.jwt import decode_token from .core.auth.api_key import hash_api_key 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: 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: 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 with optional server-side filtering. 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 stream connection", extra={ "task_id": task_id, "source_filter": source_filter, "level_filter": level_filter, "min_level": min_level, }, ) task_manager = get_task_manager() queue = await task_manager.subscribe_logs(task_id) logger.reason( "Subscribed WebSocket client to task log queue", 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 try: 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.dict() 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), }, ) 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}, ) while True: log_entry = await queue.get() if not matches_filters(log_entry): continue log_dict = log_entry.dict() 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 log_entry.message or "Task failed" in log_entry.message ): logger.reason( "Observed terminal task log entry; delaying to preserve client visibility", extra={"task_id": task_id, "message": log_entry.message}, ) await asyncio.sleep(2) except WebSocketDisconnect: logger.reason( "WebSocket client disconnected from task log stream", extra={"task_id": task_id}, ) except Exception as exc: logger.explore( "WebSocket log streaming encountered an unexpected failure", extra={"task_id": task_id, "error": str(exc)}, ) raise finally: task_manager.unsubscribe_logs(task_id, queue) logger.reflect( "Released WebSocket log queue subscription", extra={"task_id": task_id}, ) # #endregion websocket_endpoint # #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 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