Spaces:
Sleeping
Sleeping
| """FastAPI application entry point.""" | |
| from __future__ import annotations | |
| import os | |
| from contextlib import asynccontextmanager | |
| from pathlib import Path | |
| from fastapi import FastAPI, Depends, Request, HTTPException | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.staticfiles import StaticFiles | |
| from fastapi.responses import FileResponse, JSONResponse | |
| from starlette.middleware.base import BaseHTTPMiddleware | |
| from visual_notebook.config import settings | |
| from visual_notebook.session import Session, SessionStore | |
| from visual_notebook.api.routes import router as api_router | |
| class CacheControlMiddleware(BaseHTTPMiddleware): | |
| """Set proper Cache-Control headers so new deploys are immediately visible. | |
| - index.html: no-cache (always revalidate with server) | |
| - /assets/* (Vite hashed filenames): long cache + immutable | |
| - /api/* and /health: no-store (never cache API responses) | |
| - everything else: no-cache | |
| """ | |
| async def dispatch(self, request, call_next): | |
| response = await call_next(request) | |
| path = request.url.path | |
| if path.startswith("/api/") or path == "/health": | |
| response.headers["Cache-Control"] = "no-store, no-cache, must-revalidate" | |
| response.headers["Pragma"] = "no-cache" | |
| response.headers["Expires"] = "0" | |
| elif path == "/" or path == "/index.html": | |
| response.headers["Cache-Control"] = "no-cache, no-store, must-revalidate" | |
| response.headers["Pragma"] = "no-cache" | |
| response.headers["Expires"] = "0" | |
| elif "/assets/" in path: | |
| # Vite generates hashed filenames — safe to cache aggressively | |
| response.headers["Cache-Control"] = "public, max-age=31536000, immutable" | |
| elif path.startswith("/static/"): | |
| response.headers["Cache-Control"] = "no-cache" | |
| else: | |
| response.headers["Cache-Control"] = "no-cache" | |
| return response | |
| async def lifespan(app: FastAPI): | |
| """Startup/shutdown lifecycle.""" | |
| Path(settings.data_dir).mkdir(parents=True, exist_ok=True) | |
| Path(settings.checkpoint_dir).mkdir(parents=True, exist_ok=True) | |
| app.state.session_store = SessionStore() | |
| app.state.demo_loaded = False | |
| app.state.demo_table = None | |
| app.state.demo_row_count = 0 | |
| app.state.demo_columns = [] | |
| app.state.demo_filename = None | |
| # in-memory to avoid lock file issues | |
| import duckdb | |
| app.state.duckdb_con = duckdb.connect(":memory:", read_only=False) | |
| # Auto-load demo data on startup | |
| _autoload_demo(app) | |
| yield | |
| app.state.duckdb_con.close() | |
| def _autoload_demo(app: FastAPI) -> None: | |
| """Auto-load the first CSV found in DATA_DIR as a shared demo dataset. | |
| This ensures the app isn't empty when deployed to HF Spaces or Docker. | |
| The demo is registered in DuckDB at startup and any session can load it. | |
| """ | |
| data_dir = Path(settings.data_dir) | |
| csv_files = sorted(data_dir.glob("*.csv")) | |
| if not csv_files: | |
| return | |
| from visual_notebook.data.csv_loader import register_csv | |
| from visual_notebook.data.schema import infer_schema_duckdb | |
| con = app.state.duckdb_con | |
| demo_path = csv_files[0] | |
| try: | |
| meta = register_csv(con, "demo_dataset", str(demo_path)) | |
| schema = infer_schema_duckdb(con, "demo_dataset") | |
| app.state.demo_loaded = True | |
| app.state.demo_table = "demo_dataset" | |
| app.state.demo_row_count = meta["row_count"] | |
| app.state.demo_columns = schema | |
| app.state.demo_filename = demo_path.name | |
| print(f"[startup] Auto-loaded demo dataset: {demo_path.name} ({meta['row_count']} rows)") | |
| except Exception as e: | |
| print(f"[startup] Failed to auto-load demo: {e}") | |
| app = FastAPI( | |
| title="Visual Notebook", | |
| version="0.2.0", | |
| lifespan=lifespan, | |
| ) | |
| # Cache control middleware MUST be added before CORS so headers are set early | |
| app.add_middleware(CacheControlMiddleware) | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=settings.cors_origins, | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| app.include_router(api_router, prefix="/api") | |
| async def health(): | |
| return {"status": "ok", "version": "0.2.0"} | |
| async def get_demo_info(request: Request): | |
| """Return demo dataset info so the frontend can offer to load it.""" | |
| if not request.app.state.demo_loaded: | |
| return {"available": False} | |
| from visual_notebook.api.schemas import ColumnInfo | |
| return { | |
| "available": True, | |
| "filename": request.app.state.demo_filename, | |
| "row_count": request.app.state.demo_row_count, | |
| "columns": [ColumnInfo(**c) for c in request.app.state.demo_columns], | |
| } | |
| async def load_demo(request: Request): | |
| """Load the demo dataset into a new or existing session.""" | |
| if not request.app.state.demo_loaded: | |
| raise HTTPException(status_code=404, detail="No demo dataset available") | |
| from visual_notebook.domain.dataframe import DataFrameRef | |
| from visual_notebook.domain.history import History | |
| from visual_notebook.data.schema import infer_schema_duckdb | |
| from visual_notebook.api.schemas import ColumnInfo, UploadResponse | |
| con = request.app.state.duckdb_con | |
| store = request.app.state.session_store | |
| session_id = request.cookies.get("session_id") | |
| session = store.get_or_create(session_id, con) | |
| table_name = request.app.state.demo_table | |
| ref = DataFrameRef( | |
| backend="duckdb", | |
| name=table_name, | |
| con=con, | |
| source_path=str(Path(settings.data_dir) / request.app.state.demo_filename), | |
| ) | |
| session.current_table = table_name | |
| session.current_ref = ref | |
| session.history = History() | |
| session.history.commit(ref, operation="load_demo", params={"filename": request.app.state.demo_filename}) | |
| schema = infer_schema_duckdb(con, table_name) | |
| response = UploadResponse( | |
| dataset_id="demo", | |
| table_name=table_name, | |
| row_count=request.app.state.demo_row_count, | |
| columns=[ColumnInfo(**c) for c in schema], | |
| session_id=session.id, | |
| ) | |
| resp = JSONResponse(content=response.model_dump()) | |
| resp.set_cookie( | |
| "session_id", | |
| session.id, | |
| httponly=True, | |
| max_age=86400 * 7, | |
| samesite="none", | |
| secure=True, | |
| ) | |
| return resp | |
| # Static frontend serving (for production / Docker / HF Spaces) | |
| STATIC_DIR = Path(os.getenv("STATIC_DIR", "/app/static")) | |
| def _try_mount_static() -> None: | |
| """Attempt to mount static frontend files. No-op if directory doesn't exist.""" | |
| try: | |
| if not STATIC_DIR.is_dir(): | |
| return | |
| except (PermissionError, OSError): | |
| return | |
| assets_dir = STATIC_DIR / "assets" | |
| try: | |
| if assets_dir.is_dir(): | |
| app.mount("/assets", StaticFiles(directory=str(assets_dir)), name="static-assets") | |
| except (PermissionError, OSError): | |
| pass | |
| try: | |
| app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static-root") | |
| except (PermissionError, OSError): | |
| pass | |
| async def serve_spa(full_path: str): | |
| """Serve the React SPA - fall back to index.html for client-side routing.""" | |
| try: | |
| file_path = STATIC_DIR / full_path | |
| if file_path.is_file(): | |
| return FileResponse(str(file_path)) | |
| except (PermissionError, OSError): | |
| pass | |
| return FileResponse(str(STATIC_DIR / "index.html")) | |
| _try_mount_static() | |