diff --git a/mini-program/app/db.py b/mini-program/app/db.py new file mode 100644 index 0000000..d07fa46 --- /dev/null +++ b/mini-program/app/db.py @@ -0,0 +1,45 @@ +from collections.abc import Generator + +from sqlalchemy import create_engine, text +from sqlalchemy.orm import Session, sessionmaker + +try: + # For module mode: `uvicorn app.main:app` + from app.settings import settings +except ModuleNotFoundError: + # For script mode: `python app/main.py` or VS Code "Run Python File" + from settings import settings + + +engine = create_engine( + settings.mysql_dsn, + pool_size=5, + max_overflow=10, + pool_pre_ping=True, + pool_recycle=1800, + future=True, +) + +SessionLocal = sessionmaker( + bind=engine, + autocommit=False, + autoflush=False, + expire_on_commit=False, +) + + +def get_db() -> Generator[Session, None, None]: + db = SessionLocal() + try: + yield db + finally: + db.close() + + +def check_db_connection() -> None: + with engine.connect() as conn: + conn.execute(text("SELECT 1")) + + +def close_db_engine() -> None: + engine.dispose() diff --git a/mini-program/app/main.py b/mini-program/app/main.py index 60084ed..aa3a223 100644 --- a/mini-program/app/main.py +++ b/mini-program/app/main.py @@ -4,10 +4,12 @@ from fastapi import FastAPI try: # For module mode: `uvicorn app.main:app` + from app.db import check_db_connection, close_db_engine from app.routers.health import router as health_router from app.settings import settings except ModuleNotFoundError: # For script mode: `python app/main.py` or VS Code "Run Python File" + from db import check_db_connection, close_db_engine from routers.health import router as health_router from settings import settings @@ -18,6 +20,15 @@ def create_app() -> FastAPI: version=settings.app_version, description=settings.app_description, ) + + @app.on_event("startup") + def on_startup() -> None: + check_db_connection() + + @app.on_event("shutdown") + def on_shutdown() -> None: + close_db_engine() + app.include_router(health_router) return app diff --git a/mini-program/requirements.txt b/mini-program/requirements.txt index 3f5d45a..4df36d6 100644 --- a/mini-program/requirements.txt +++ b/mini-program/requirements.txt @@ -2,3 +2,5 @@ fastapi==0.115.11 uvicorn[standard]==0.34.0 pydantic-settings==2.8.1 python-dotenv==1.0.1 +SQLAlchemy==2.0.38 +PyMySQL==1.1.1