- Simplified init_db() to only call Base.metadata.create_all() on all
three engines — no more _ensure_*() inline additive migrations
- run_alembic_migrations() now raises on failure (no safety net to
fall back to)
- All schema changes (add/drop/rename columns, data migrations) must
go through Alembic. create_all() handles new tables only.
@REJECTED _ensure_*() inline migrations removed because they:
- Duplicated Alembic logic without versioning
- Created hidden schema drift (columns existed in DB but had no
corresponding Alembic migration)
- Made audits and fresh DB provisioning unreliable
628 lines
27 KiB
Python
Executable File
628 lines
27 KiB
Python
Executable File
# #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
|