import os import logging import asyncio import tempfile import shutil import uuid from contextlib import asynccontextmanager from datetime import datetime, timezone from fastapi import FastAPI, UploadFile, File, HTTPException, Depends, Request, Response, status from fastapi.middleware.cors import CORSMiddleware from fastapi.staticfiles import StaticFiles from sqlalchemy import text from sqlalchemy.orm import Session from app.middleware.csrf import CSRFMiddleware from app.middleware.ip_middleware import IPMiddleware from app.core.settings import settings from app.db.database import init_db, init_db_pool, close_db_pool, get_db import app.db.all_models # noqa: F401 from app.middleware.auth import get_current_user from app.modules.auth.models.user_model import User from app.modules.auth.models.access_model import Access from app.modules.auth.models.role_model import Role from app.modules.auth.models.role_access_model import RoleAccess from app.modules.documents.models.document_model import Project from app.modules.auth.routes.auth_routes import router as auth_router from app.modules.auth.routes.user_routes import router as user_router from app.modules.auth.routes.me_routes import router as me_router from app.modules.auth.routes.access_routes import router as access_router from app.modules.auth.routes.role_routes import router as role_router from app.modules.auth.routes.admin_dashboard_routes import router as admin_dashboard_router from app.modules.auth.routes.admin_user_routes import router as admin_user_router from app.modules.documents.routes.document_routes import router as projects_router from app.modules.documents.routes.sharing_routes import router as sharing_router from app.modules.documents.routes.version_routes import router as version_router from app.modules.documents.routes.comment_routes import ( router as project_comments_router, ) from app.modules.documents.routes.details_routes import router as project_details_router from app.modules.drive.routes.drive_routes import ( router as drive_router, public_router as public_drive_router, ) from app.modules.storage.routes.storage_routes import router as storage_router from app.modules.chat.routes.chat_routes import router as chatbot_router from app.tasks.routes import router as task_results_router from app.modules.notifications.routes.notification_routes import ( router as notifications_router, ) from app.modules.documents.routes.export_routes import router as export_router from app.modules.speech.routes.speech_routes import router as speech_router from app.api.ws_router import router as ws_router from app.core.pubsub_listener import start_listener from app.core.redis import redis_pubsub from app.modules.collab.routes.chat_routes import ( router as collab_chat_router, rest_router as collab_rest_router, ) from app.modules.tenant.routes.tenant_routes import router as tenant_router from app.modules.tenant.routes.tenant_contact_routes import router as tenant_contact_router from app.modules.signing.routes.signing_routes import router as signing_router from app.modules.signing.routes.config_routes import router as signing_config_router from app.modules.configuration.routes.system_configuration_routes import ( router as system_configuration_router, ) from app.modules.activity_logs.routes import router as activity_logs_router from app.infrastructure.realtime.chat_subscriber import ( start_chat_subscriber, stop_chat_subscriber, ) from app.modules.drive.routes.scan_routes import router as scan_router from app.modules.extraction.routes.extraction_routes import router as extraction_router from app.modules.editor.pdf_router import router as pdf_editor_router from app.modules.auth.sso_router import router as sso_router from app.modules.security.routes.security_routes import router as security_router TEMP_DIR = tempfile.mkdtemp(prefix="ai_uploads_") processing_semaphore = asyncio.Semaphore(1) logging.basicConfig( level=getattr(logging, settings.LOG_LEVEL), format="%(asctime)s - %(levelname)s - %(message)s", ) @asynccontextmanager async def lifespan(app: FastAPI): logging.info(f"Server is starting up in {settings.APP_ENV} mode...") try: init_db_pool() init_db() asyncio.create_task(start_listener()) await redis_pubsub.connect() await start_chat_subscriber() from app.infrastructure.vector.vector_client import preload as preload_vector loop = asyncio.get_running_loop() loop.run_in_executor(None, preload_vector) except Exception as e: logging.error(f"Critical Startup failure: {e}") raise e yield logging.info("Server is shutting down...") try: await stop_chat_subscriber() await redis_pubsub.disconnect() logging.info("โœ… Chat subscriber & Redis Pub/Sub disconnected") except Exception as e: logging.error(f"Shutdown error during subscriber/redis stop: {e}") try: close_db_pool() logging.info("โœ… Database connection pool closed") except Exception as e: logging.error(f"Error closing database pool: {e}") try: shutil.rmtree(TEMP_DIR, ignore_errors=True) logging.info("๐Ÿงน Temp directory cleaned") except Exception as e: logging.error(f"Shutdown error during temp cleanup: {e}") app = FastAPI( lifespan=lifespan, title="DocQube Document Processing Pipeline API", description="An interactive API and UI to process PDFs into JATS XML.", version="2.0.0", ) app.add_middleware(CSRFMiddleware) app.add_middleware(IPMiddleware) from app.middleware.tenant_context_middleware import TenantContextMiddleware # noqa: E402 app.add_middleware(TenantContextMiddleware) app.add_middleware( CORSMiddleware, allow_origins=settings.CORS_ORIGINS_LIST, allow_credentials=True, allow_methods=["*"], allow_headers=["*"], # Browsers only expose a small header safelist to JS on cross-origin # responses by default (Content-Range/Accept-Ranges/Content-Length are # NOT in it) โ€” pdf.js reads these itself to decide whether it can use # HTTP Range requests for progressive PDF loading, so without this it # falls back to downloading the entire file before rendering anything. expose_headers=["Content-Range", "Accept-Ranges", "Content-Length"], ) # File Serving Internal Helper async def _get_s3_file_response(key: str, db: Session = None, tenant_id: str = None): from app.infrastructure.storage.local_storage_handler import ( get_storage_client, ) from fastapi.responses import StreamingResponse import mimetypes storage, quarantine_bucket, clean_bucket = get_storage_client( db=db, tenant_id=tenant_id) try: try: response = storage.get_object(Bucket=clean_bucket, Key=key) except Exception: response = storage.get_object(Bucket=quarantine_bucket, Key=key) body = response["Body"] content_type, _ = mimetypes.guess_type(key) if not content_type: content_type = "application/octet-stream" return StreamingResponse(body, media_type=content_type) except HTTPException: raise except Exception as e: logging.error(f"Storage File Proxy Error: {e}") raise HTTPException(status_code=404, detail="File not found") @app.get("/api/storage_drive/{key:path}") async def serve_storage_file(key: str, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)): if ".." in key: raise HTTPException(status_code=400, detail="Invalid path") tenant_part = f"tenants/{current_user.tenant_id}/" if current_user.tenant_id else "" user_prefix1 = f"{tenant_part}user_{current_user.id}/" user_prefix2 = f"{tenant_part}users/{current_user.id}/" conv_prefix = f"{tenant_part}doc_conversion/{current_user.id}/" chatbot_prefix = f"{tenant_part}chatbot_documents/{current_user.id}/" images_prefix = f"extracted_images/{current_user.id}/" allowed_prefixes = [ user_prefix1, user_prefix2, conv_prefix, chatbot_prefix, images_prefix, f"user_{current_user.id}/", f"users/{current_user.id}/", f"doc_conversion/{current_user.id}/", f"chatbot_documents/{current_user.id}/", f"extracted_images/{current_user.id}/", ] if not any(key.startswith(p) for p in allowed_prefixes): raise HTTPException( status_code=403, detail="Access denied to this storage key") return await _get_s3_file_response(key, db=db, tenant_id=current_user.tenant_id) @app.get("/api/serve-image/{session_id}/{filename}") async def serve_image( session_id: str, filename: str, current_user: User = Depends(get_current_user), db: Session = Depends(get_db), ): """ Alias for serving images extracted during document processing. """ if ".." in session_id or ".." in filename: raise HTTPException(status_code=400, detail="Invalid parameters") project = ( db.query(Project) .filter(Project.session_id == session_id, Project.user_id == current_user.id) .first() ) if not project: raise HTTPException( status_code=403, detail="Access denied to this session's resources" ) key = f"extracted_images/{current_user.id}/{session_id}/{filename}" return await _get_s3_file_response(key, db=db, tenant_id=current_user.tenant_id) def _check_database_ready(db: Session) -> bool: """Lightweight DB reachability check shared by /api/health/ready and the legacy /api/health. On failure it rolls back before returning โ€” without this, the session handed back to Depends(get_db)'s own cleanup (which unconditionally calls db.commit() after a successful return from this function) is left in an aborted-transaction state, and that commit() itself raises. That turns a DB failure this function already caught into an unhandled 500 from the dependency's own teardown โ€” which is exactly what took the original /api/health down during the PgBouncer incident, even though it appeared to catch the error internally. """ try: db.execute(text("SELECT 1")) return True except Exception: db.rollback() return False @app.get("/api/health/live") def health_live(): """Liveness: is the API process able to serve HTTP at all. No database, no PgBouncer, no external dependency of any kind โ€” this is what the Docker HEALTHCHECK now uses, so a temporarily unreachable database can no longer get the whole API container killed as unhealthy. """ return { "status": "ok", "service": "api", "timestamp": datetime.now(timezone.utc).isoformat(), } @app.get("/api/health/ready") def health_ready(response: Response, db: Session = Depends(get_db)): """Readiness: is the API ready to serve real traffic, i.e. can it reach the database through the application's actual configured pool/engine (app.db.database.engine via get_db) โ€” the same one every request uses, not a separate health-only connection. """ if _check_database_ready(db): return { "status": "ready", "database": "connected", "timestamp": datetime.now(timezone.utc).isoformat(), } response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE return { "status": "not_ready", "database": "disconnected", "timestamp": datetime.now(timezone.utc).isoformat(), } @app.get("/api/health") def health_check(db: Session = Depends(get_db)): """Backward-compatible general health endpoint โ€” response shape kept identical to before this refactor. Reuses the same DB-reachability check as /api/health/ready instead of duplicating it; the Docker HEALTHCHECK no longer depends on this endpoint (see /api/health/live). """ if not _check_database_ready(db): return {"status": "degraded", "database": "error", "error": "database unreachable"} health = {"status": "ok", "database": "connected"} try: version = db.execute(text("SELECT version_num FROM alembic_version")).scalar() health["migration_version"] = version except Exception: db.rollback() health["status"] = "degraded" health["error"] = "migration version lookup failed" return health app.include_router(tenant_router, prefix="/api") app.include_router(tenant_contact_router, prefix="/api/tenant") app.include_router(auth_router, prefix="/api") app.include_router(sso_router, prefix="/api") app.include_router(me_router, prefix="/api") app.include_router(access_router) app.include_router(role_router) app.include_router(admin_dashboard_router) app.include_router(admin_user_router) app.include_router(user_router, prefix="/api") app.include_router(projects_router, prefix="/api") app.include_router(sharing_router, prefix="/api") app.include_router(version_router, prefix="/api") app.include_router(project_comments_router, prefix="/api") app.include_router(project_details_router, prefix="/api") app.include_router(drive_router, prefix="/api") app.include_router(public_drive_router, prefix="/api") app.include_router(storage_router, prefix="/api") app.include_router(task_results_router, prefix="/api") app.include_router(chatbot_router) app.include_router(notifications_router, prefix="/api") app.include_router(system_configuration_router) app.include_router(activity_logs_router) app.include_router(export_router, prefix="/api") app.include_router(speech_router) app.include_router(ws_router, prefix="/api") app.include_router(collab_chat_router, prefix="/api") app.include_router(collab_rest_router, prefix="/api") app.include_router(signing_router, prefix="/api") app.include_router(signing_config_router, prefix="/api") app.include_router(pdf_editor_router, prefix="/api") app.include_router(scan_router, prefix="/api") app.include_router(extraction_router, prefix="/api") from app.modules.org.routes.org_routes import router as org_router # noqa: E402 app.include_router(org_router, prefix="/api") from app.modules.billing.routes.billing_routes import router as billing_router # noqa: E402 app.include_router(billing_router, prefix="/api") app.include_router(security_router, prefix="/api") @app.post("/api/extract") async def extract( request: Request, file: UploadFile = File(...), current_user: User = Depends(get_current_user), ): """ On-demand document extraction endpoint. Loads models only when called and unloads after. """ file_id = str(uuid.uuid4()) ext = os.path.splitext(file.filename)[1].lower() temp_path = os.path.join(TEMP_DIR, f"{file_id}{ext}") MAX_FILE_SIZE = 50 * 1024 * 1024 try: total_size = 0 with open(temp_path, "wb") as buffer: while True: chunk = await file.read(1024 * 1024) if not chunk: break total_size += len(chunk) if total_size > MAX_FILE_SIZE: raise HTTPException( status_code=413, detail="File too large (Max 50MB)" ) buffer.write(chunk) # Convert PDF to HTML locally, then derive markdown async with processing_semaphore: from app.modules.editor.converters import process_pdf_to_html loop = asyncio.get_running_loop() def _convert(): html = process_pdf_to_html(temp_path) from bs4 import BeautifulSoup markdown = BeautifulSoup( html, "html.parser").get_text("\n") return {"html": html, "markdown": markdown} result = await loop.run_in_executor(None, _convert) return { "status": "success", "file_id": file_id, "data": result, } except HTTPException: raise except Exception as e: error_id = str(uuid.uuid4()) logging.error( f"Extraction failed [Ref: {error_id}]: {e}", exc_info=True) raise HTTPException( status_code=500, detail=f"An internal error occurred during processing. Please contact support and provide reference: {error_id}", ) finally: if os.path.exists(temp_path): try: os.remove(temp_path) except Exception as e: logging.error(f"Cleanup error removing {temp_path}: {e}") if __name__ == "__main__": import uvicorn uvicorn.run( "main:app", host=settings.HOST, port=settings.PORT, reload=settings.DEBUG )