image_management_api/main.py

import logging
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi.openapi.docs import get_swagger_ui_html
from fastapi.openapi.utils import get_openapi
import time

# Import API routers
from src.api.v1 import teams, users, images, auth, search

# Import configuration and database
from src.config.config import settings
from src.utils.logging import setup_logging
from src.db import db
from src.db.providers.firestore_provider import firestore_db
from src.db.repositories import init_repositories

# Setup logging
setup_logging()
logger = logging.getLogger(__name__)

# Create FastAPI app
app = FastAPI(
    title=settings.PROJECT_NAME,
    description="API for securely storing, organizing, and retrieving images",
    version="1.0.0",
    docs_url=None,
    redoc_url=None,
    openapi_url="/api/v1/openapi.json"
)

# Add request logging middleware
@app.middleware("http")
async def log_requests(request: Request, call_next):
    start_time = time.time()
    
    # Log incoming request details
    logger.info(f"Incoming request: {request.method} {request.url}")
    logger.info(f"Headers: {dict(request.headers)}")
    logger.info(f"Client: {request.client}")
    
    response = await call_next(request)
    
    process_time = time.time() - start_time
    logger.info(f"Request completed in {process_time:.4f}s with status {response.status_code}")
    
    return response

# Connect to database
try:
    db.connect_to_database()
    logger.info("Database connection initialized")
    
    # Initialize the Firestore provider
    firestore_db.connect()
    logger.info("Firestore provider initialized")
    
    # Initialize repositories
    init_repositories()
    logger.info("Repositories initialized")
except Exception as e:
    logger.error(f"Failed to connect to database: {e}", exc_info=True)
    # We'll continue without database for Swagger UI to work, but operations will fail

# Set up CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=settings.CORS_ORIGINS,
    allow_credentials=True,
    allow_methods=settings.CORS_METHODS,
    allow_headers=settings.CORS_HEADERS,
    expose_headers=settings.CORS_EXPOSE_HEADERS,
    max_age=settings.CORS_MAX_AGE,
)

# Include API routers
app.include_router(auth.router, prefix="/api/v1")
app.include_router(teams.router, prefix="/api/v1")
app.include_router(users.router, prefix="/api/v1")
app.include_router(images.router, prefix="/api/v1")
app.include_router(search.router, prefix="/api/v1")

# Custom exception handler
@app.exception_handler(Exception)
async def general_exception_handler(request: Request, exc: Exception):
    logger.error(f"Unhandled exception: {exc}", exc_info=True)
    # Include more details for easier debugging
    error_details = {
        "detail": "Internal server error",
        "type": type(exc).__name__,
        "path": request.url.path
    }
    # Only include exception message in development mode
    if settings.ENVIRONMENT == "development":
        error_details["message"] = str(exc)
    
    return JSONResponse(
        status_code=500,
        content=error_details
    )

# Custom Swagger UI with API key authentication
@app.get("/docs", include_in_schema=False)
async def custom_swagger_ui_html():
    return get_swagger_ui_html(
        openapi_url=app.openapi_url,
        title=f"{app.title} - Swagger UI",
        oauth2_redirect_url=app.swagger_ui_oauth2_redirect_url,
        swagger_js_url="https://cdn.jsdelivr.net/npm/swagger-ui-dist@5.9.0/swagger-ui-bundle.js",
        swagger_css_url="https://cdn.jsdelivr.net/npm/swagger-ui-dist@5.9.0/swagger-ui.css",
    )

# Custom OpenAPI schema to include API key auth
def custom_openapi():
    if app.openapi_schema:
        return app.openapi_schema
    
    openapi_schema = get_openapi(
        title=app.title,
        version=app.version,
        description=app.description,
        routes=app.routes,
    )
    
    # Initialize components if not present
    if "components" not in openapi_schema:
        openapi_schema["components"] = {}
    
    # Add API key security scheme
    openapi_schema["components"]["securitySchemes"] = {
        "ApiKeyAuth": {
            "type": "apiKey",
            "in": "header",
            "name": "X-API-Key"
        }
    }
    
    # Ensure schemas are included
    if "schemas" not in openapi_schema["components"]:
        openapi_schema["components"]["schemas"] = {}
    
    # Apply security to endpoints except auth, users, teams, and API key creation endpoints
    for path in openapi_schema["paths"]:
        # Exclude auth endpoints, users endpoints, teams endpoints, and API key creation
        if not (path.startswith("/api/v1/auth") or 
                path.startswith("/api/v1/users") or 
                path.startswith("/api/v1/teams")):
            for method in openapi_schema["paths"][path]:
                if method.lower() in ["get", "post", "put", "delete", "patch"]:
                    if "security" not in openapi_schema["paths"][path][method]:
                        openapi_schema["paths"][path][method]["security"] = [{"ApiKeyAuth": []}]
    
    app.openapi_schema = openapi_schema
    return app.openapi_schema

app.openapi = custom_openapi

@app.get("/health", include_in_schema=False)
async def health_check():
    """Health check endpoint for monitoring and client connectivity tests"""
    return {"status": "healthy", "message": "API is running"}

@app.get("/", include_in_schema=False)
async def root():
    return {"message": "Welcome to the Image Management API. Please see /docs for API documentation."}

# Shutdown handler to close database connections
@app.on_event("shutdown")
async def shutdown_event():
    logger.info("Application shutting down")
    db.close_database_connection()
    # Also disconnect the Firestore provider
    firestore_db.disconnect()
    logger.info("Firestore provider disconnected")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)