Unified_Authentication_Platform/backend/app/main.py

import logging
from contextlib import asynccontextmanager

from fastapi import FastAPI
from starlette.middleware.cors import CORSMiddleware
from tenacity import retry, stop_after_attempt, wait_fixed, before_log, after_log

from app.api.v1.api import api_router
from app.core.config import settings
from app.core.database import Base, engine
from app.core.cache import redis_client
from app.core.logging_config import setup_logging
from app.core.scheduler import scheduler
from app.services.backup_service import BackupService
from app.services.db_backup_upload_service import DatabaseBackupUploadService
from app.core.init_data import fix_user_names
from app.core.database import SessionLocal

# Configure Logging
logger = setup_logging()

# Retry Logic Configuration
max_tries = 60 * 5  # 5 minutes
wait_seconds = 1

@retry(
    stop=stop_after_attempt(max_tries),
    wait=wait_fixed(wait_seconds),
    before=before_log(logger, logging.INFO),
    after=after_log(logger, logging.WARN),
)
def init_db() -> None:
    try:
        # Try to create a connection to check if DB is ready
        with engine.connect() as conn:
            pass
        # Create tables
        Base.metadata.create_all(bind=engine)
        logger.info("Database connection established and tables created.")
    except Exception as e:
        logger.error(e)
        raise e

@retry(
    stop=stop_after_attempt(max_tries),
    wait=wait_fixed(wait_seconds),
    before=before_log(logger, logging.INFO),
    after=after_log(logger, logging.WARN),
)
def check_redis() -> None:
    try:
        redis_client.ping()
        logger.info("Redis connection established.")
    except Exception as e:
        logger.error(e)
        raise e

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup: Wait for services
    logger.info("Waiting for Database and Redis...")
    init_db()
    check_redis()
    
    # Init Scheduler
    logger.info("Starting Scheduler...")
    scheduler.start()
    try:
        with SessionLocal() as db:
            # Check and fix user names/english_names
            fix_user_names(db)
            
            BackupService.init_scheduler(db)
            
            # 初始化数据库备份上传服务
            DatabaseBackupUploadService.init_scheduler()
    except Exception as e:
        logger.error(f"Failed to init scheduler or fix user names: {e}")

    yield
    
    # Shutdown logic
    logger.info("Shutting down WebSocket manager...")
    from app.core.websocket_manager import manager
    await manager.shutdown()
    
    scheduler.shutdown()
    logger.info("Shutting down...")

app = FastAPI(
    title=settings.PROJECT_NAME,
    openapi_url=f"{settings.API_V1_STR}/openapi.json",
    docs_url=f"{settings.API_V1_STR}/docs",
    redoc_url=f"{settings.API_V1_STR}/redoc",
    lifespan=lifespan
)

# Allow all CORS for development simplicity if config is restrictive
origins = [
    "http://localhost:5173",
    "http://127.0.0.1:5173",
    "http://localhost:8000",
    "http://127.0.0.1:8000"
]

# Set all CORS enabled origins
if settings.BACKEND_CORS_ORIGINS:
    for origin in settings.BACKEND_CORS_ORIGINS:
        origins.append(str(origin))

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["X-New-Token"],
)

app.include_router(api_router, prefix=settings.API_V1_STR)

@app.get("/")
def root():
    return {"message": "Welcome to Unified Authentication Platform API"}