feat: Add SaaS robustness middleware - Rate limiting with token bucket and sliding window algorithms - Input validation (file, language, provider) - Security headers middleware (CSP, XSS protection, etc.) - Automatic file cleanup with TTL tracking - Memory and disk monitoring - Enhanced health check and metrics endpoints - Request logging with unique IDs

main.py

@@ -1,24 +1,55 @@
 """
 Document Translation API
 FastAPI application for translating complex documents while preserving formatting
+SaaS-ready with rate limiting, validation, and robust error handling
 """
-from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Request, Depends
 from fastapi.responses import FileResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
+from contextlib import asynccontextmanager
 from pathlib import Path
 from typing import Optional
 import asyncio
 import logging
+import os
 
 from config import config
 from translators import excel_translator, word_translator, pptx_translator
 from utils import file_handler, handle_translation_error, DocumentProcessingError
 
-# Configure logging
-logging.basicConfig(level=logging.INFO)
+# Import SaaS middleware
+from middleware.rate_limiting import RateLimitMiddleware, RateLimitManager, RateLimitConfig
+from middleware.security import SecurityHeadersMiddleware, RequestLoggingMiddleware, ErrorHandlingMiddleware
+from middleware.cleanup import FileCleanupManager, MemoryMonitor, HealthChecker, create_cleanup_manager
+from middleware.validation import FileValidator, LanguageValidator, ProviderValidator, InputSanitizer, ValidationError
+
+# Configure structured logging
+logging.basicConfig(
+    level=getattr(logging, os.getenv("LOG_LEVEL", "INFO")),
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
 logger = logging.getLogger(__name__)
 
+# Initialize SaaS components
+rate_limit_config = RateLimitConfig(
+    requests_per_minute=int(os.getenv("RATE_LIMIT_PER_MINUTE", "30")),
+    requests_per_hour=int(os.getenv("RATE_LIMIT_PER_HOUR", "200")),
+    translations_per_minute=int(os.getenv("TRANSLATIONS_PER_MINUTE", "10")),
+    translations_per_hour=int(os.getenv("TRANSLATIONS_PER_HOUR", "50")),
+    max_concurrent_translations=int(os.getenv("MAX_CONCURRENT_TRANSLATIONS", "5")),
+)
+rate_limit_manager = RateLimitManager(rate_limit_config)
+
+cleanup_manager = create_cleanup_manager(config)
+memory_monitor = MemoryMonitor(max_memory_percent=float(os.getenv("MAX_MEMORY_PERCENT", "80")))
+health_checker = HealthChecker(cleanup_manager, memory_monitor)
+
+file_validator = FileValidator(
+    max_size_mb=config.MAX_FILE_SIZE_MB,
+    allowed_extensions=config.SUPPORTED_EXTENSIONS
+)
+
 
 def build_full_prompt(system_prompt: str, glossary: str) -> str:
     """Combine system prompt and glossary into a single prompt for LLM translation."""
@@ -40,23 +71,47 @@ Always use the translations from this glossary when you encounter these terms.""
     return "\n\n".join(parts) if parts else ""
 
 
-# Ensure necessary directories exist
-config.ensure_directories()
+# Lifespan context manager for startup/shutdown
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Handle startup and shutdown events"""
+    # Startup
+    logger.info("Starting Document Translation API...")
+    config.ensure_directories()
+    await cleanup_manager.start()
+    logger.info("API ready to accept requests")
+    
+    yield
+    
+    # Shutdown
+    logger.info("Shutting down...")
+    await cleanup_manager.stop()
+    logger.info("Cleanup completed")
 
-# Create FastAPI app
+
+# Create FastAPI app with lifespan
 app = FastAPI(
     title=config.API_TITLE,
     version=config.API_VERSION,
-    description=config.API_DESCRIPTION
+    description=config.API_DESCRIPTION,
+    lifespan=lifespan
 )
 
-# Add CORS middleware
+# Add middleware (order matters - first added is outermost)
+app.add_middleware(ErrorHandlingMiddleware)
+app.add_middleware(RequestLoggingMiddleware, log_body=False)
+app.add_middleware(SecurityHeadersMiddleware, config={"enable_hsts": os.getenv("ENABLE_HSTS", "false").lower() == "true"})
+app.add_middleware(RateLimitMiddleware, rate_limit_manager=rate_limit_manager)
+
+# CORS - configure for production
+allowed_origins = os.getenv("CORS_ORIGINS", "*").split(",")
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Configure appropriately for production
+    allow_origins=allowed_origins,
     allow_credentials=True,
-    allow_methods=["*"],
+    allow_methods=["GET", "POST", "DELETE", "OPTIONS"],
     allow_headers=["*"],
+    expose_headers=["X-Request-ID", "X-Original-Filename", "X-File-Size-MB", "X-Target-Language"]
 )
 
 # Mount static files
@@ -65,6 +120,20 @@ if static_dir.exists():
     app.mount("/static", StaticFiles(directory=str(static_dir)), name="static")
 
 
+# Custom exception handler for ValidationError
+@app.exception_handler(ValidationError)
+async def validation_error_handler(request: Request, exc: ValidationError):
+    """Handle validation errors with user-friendly messages"""
+    return JSONResponse(
+        status_code=400,
+        content={
+            "error": exc.code,
+            "message": exc.message,
+            "details": exc.details
+        }
+    )
+
+
 @app.get("/")
 async def root():
     """Root endpoint with API information"""
@@ -83,11 +152,24 @@ async def root():
 
 @app.get("/health")
 async def health_check():
-    """Health check endpoint"""
-    return {
-        "status": "healthy",
-        "translation_service": config.TRANSLATION_SERVICE
-    }
+    """Health check endpoint with detailed system status"""
+    health_status = await health_checker.check_health()
+    status_code = 200 if health_status.get("status") == "healthy" else 503
+    
+    return JSONResponse(
+        status_code=status_code,
+        content={
+            "status": health_status.get("status", "unknown"),
+            "translation_service": config.TRANSLATION_SERVICE,
+            "memory": health_status.get("memory", {}),
+            "disk": health_status.get("disk", {}),
+            "cleanup_service": health_status.get("cleanup_service", {}),
+            "rate_limits": {
+                "requests_per_minute": rate_limit_config.requests_per_minute,
+                "translations_per_minute": rate_limit_config.translations_per_minute,
+            }
+        }
+    )
 
 
 @app.get("/languages")
@@ -128,6 +210,7 @@ async def get_supported_languages():
 
 @app.post("/translate")
 async def translate_document(
+    request: Request,
     file: UploadFile = File(..., description="Document file to translate (.xlsx, .docx, or .pptx)"),
     target_language: str = Form(..., description="Target language code (e.g., 'es', 'fr', 'de')"),
     source_language: str = Form(default="auto", description="Source language code (default: auto-detect)"),
@@ -160,11 +243,38 @@ async def translate_document(
     """
     input_path = None
     output_path = None
+    request_id = getattr(request.state, 'request_id', 'unknown')
     
     try:
+        # Validate inputs
+        sanitized_language = InputSanitizer.sanitize_language_code(target_language)
+        LanguageValidator.validate(sanitized_language)
+        ProviderValidator.validate(provider)
+        
+        # Validate file before processing
+        validation_result = await file_validator.validate_async(file)
+        if not validation_result.is_valid:
+            raise ValidationError(
+                message=f"File validation failed: {'; '.join(validation_result.errors)}",
+                code="INVALID_FILE",
+                details={"errors": validation_result.errors, "warnings": validation_result.warnings}
+            )
+        
+        # Log any warnings
+        if validation_result.warnings:
+            logger.warning(f"[{request_id}] File validation warnings: {validation_result.warnings}")
+        
+        # Check rate limit for translations
+        client_ip = request.client.host if request.client else "unknown"
+        if not await rate_limit_manager.check_translation_limit(client_ip):
+            raise HTTPException(
+                status_code=429,
+                detail="Translation rate limit exceeded. Please try again later."
+            )
+        
         # Validate file extension
         file_extension = file_handler.validate_file_extension(file.filename)
-        logger.info(f"Processing {file_extension} file: {file.filename}")
+        logger.info(f"[{request_id}] Processing {file_extension} file: {file.filename}")
         
         # Validate file size
         file_handler.validate_file_size(file)
@@ -178,7 +288,11 @@ async def translate_document(
         output_path = config.OUTPUT_DIR / output_filename
         
         await file_handler.save_upload_file(file, input_path)
-        logger.info(f"Saved input file to: {input_path}")
+        logger.info(f"[{request_id}] Saved input file to: {input_path}")
+        
+        # Track file for cleanup
+        await cleanup_manager.track_file(input_path, ttl_minutes=30)
+        await cleanup_manager.track_file(output_path, ttl_minutes=60)
         
         # Configure translation provider
         from services.translation_service import GoogleTranslationProvider, DeepLTranslationProvider, LibreTranslationProvider, OllamaTranslationProvider, OpenAITranslationProvider, translation_service
@@ -657,6 +771,74 @@ async def reconstruct_document(
         raise HTTPException(status_code=500, detail=f"Failed to reconstruct document: {str(e)}")
 
 
+# ============== SaaS Management Endpoints ==============
+
+@app.get("/metrics")
+async def get_metrics():
+    """Get system metrics and statistics for monitoring"""
+    health_status = await health_checker.check_health()
+    cleanup_stats = cleanup_manager.get_stats()
+    rate_limit_stats = rate_limit_manager.get_stats()
+    
+    return {
+        "system": {
+            "memory": health_status.get("memory", {}),
+            "disk": health_status.get("disk", {}),
+            "status": health_status.get("status", "unknown")
+        },
+        "cleanup": cleanup_stats,
+        "rate_limits": rate_limit_stats,
+        "config": {
+            "max_file_size_mb": config.MAX_FILE_SIZE_MB,
+            "supported_extensions": list(config.SUPPORTED_EXTENSIONS),
+            "translation_service": config.TRANSLATION_SERVICE
+        }
+    }
+
+
+@app.get("/rate-limit/status")
+async def get_rate_limit_status(request: Request):
+    """Get current rate limit status for the requesting client"""
+    client_ip = request.client.host if request.client else "unknown"
+    status = await rate_limit_manager.get_client_status(client_ip)
+    
+    return {
+        "client_ip": client_ip,
+        "limits": {
+            "requests_per_minute": rate_limit_config.requests_per_minute,
+            "requests_per_hour": rate_limit_config.requests_per_hour,
+            "translations_per_minute": rate_limit_config.translations_per_minute,
+            "translations_per_hour": rate_limit_config.translations_per_hour
+        },
+        "current_usage": status
+    }
+
+
+@app.post("/admin/cleanup/trigger")
+async def trigger_cleanup():
+    """Trigger manual cleanup of expired files"""
+    try:
+        cleaned = await cleanup_manager.cleanup_expired()
+        return {
+            "status": "success",
+            "files_cleaned": cleaned,
+            "message": f"Cleaned up {cleaned} expired files"
+        }
+    except Exception as e:
+        logger.error(f"Manual cleanup failed: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Cleanup failed: {str(e)}")
+
+
+@app.get("/admin/files/tracked")
+async def get_tracked_files():
+    """Get list of currently tracked files"""
+    tracked = cleanup_manager.get_tracked_files()
+    return {
+        "count": len(tracked),
+        "files": tracked
+    }
+
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)