128 lines
3.8 KiB
Python
128 lines
3.8 KiB
Python
from contextlib import asynccontextmanager
|
|
from typing import AsyncGenerator
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
from sqlalchemy.orm import DeclarativeBase
|
|
from sqlalchemy.pool import NullPool
|
|
|
|
from config import db_settings
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
pass
|
|
|
|
|
|
# 데이터베이스 엔진 생성
|
|
engine = create_async_engine(
|
|
url=db_settings.MYSQL_URL,
|
|
echo=False,
|
|
pool_size=10,
|
|
max_overflow=10,
|
|
pool_timeout=5,
|
|
pool_recycle=3600,
|
|
pool_pre_ping=True,
|
|
pool_reset_on_return="rollback",
|
|
connect_args={
|
|
"connect_timeout": 3,
|
|
"charset": "utf8mb4",
|
|
# "allow_public_key_retrieval": True,
|
|
},
|
|
)
|
|
|
|
# Async sessionmaker 생성
|
|
AsyncSessionLocal = async_sessionmaker(
|
|
bind=engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False,
|
|
autoflush=False, # 명시적 flush 권장
|
|
)
|
|
|
|
|
|
async def create_db_tables():
|
|
import asyncio
|
|
|
|
# 모델 import (테이블 메타데이터 등록용)
|
|
from app.home.models import Image, Project # noqa: F401
|
|
from app.lyric.models import Lyric # noqa: F401
|
|
from app.song.models import Song # noqa: F401
|
|
from app.video.models import Video # noqa: F401
|
|
|
|
print("Creating database tables...")
|
|
|
|
async with asyncio.timeout(10):
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(Base.metadata.create_all)
|
|
|
|
|
|
# FastAPI 의존성용 세션 제너레이터
|
|
async def get_session() -> AsyncGenerator[AsyncSession, None]:
|
|
async with AsyncSessionLocal() as session:
|
|
try:
|
|
yield session
|
|
# print("Session commited")
|
|
# await session.commit()
|
|
except Exception as e:
|
|
await session.rollback()
|
|
print(f"Session rollback due to: {e}")
|
|
raise e
|
|
# async with 종료 시 session.close()가 자동 호출됨
|
|
|
|
|
|
# 앱 종료 시 엔진 리소스 정리 함수
|
|
async def dispose_engine() -> None:
|
|
await engine.dispose()
|
|
print("Database engine disposed")
|
|
|
|
|
|
# =============================================================================
|
|
# 백그라운드 태스크용 세션 (별도 이벤트 루프에서 사용)
|
|
# =============================================================================
|
|
|
|
|
|
@asynccontextmanager
|
|
async def get_worker_session() -> AsyncGenerator[AsyncSession, None]:
|
|
"""백그라운드 태스크용 세션 컨텍스트 매니저
|
|
|
|
asyncio.run()으로 새 이벤트 루프를 생성하는 백그라운드 태스크에서 사용합니다.
|
|
NullPool을 사용하여 연결 풀링을 비활성화하고, 이벤트 루프 충돌을 방지합니다.
|
|
|
|
get_session()과의 차이점:
|
|
- get_session(): FastAPI DI용, 메인 이벤트 루프의 연결 풀 사용
|
|
- get_worker_session(): 백그라운드 태스크용, NullPool로 매번 새 연결 생성
|
|
|
|
Usage:
|
|
async with get_worker_session() as session:
|
|
result = await session.execute(select(Model))
|
|
await session.commit()
|
|
|
|
Note:
|
|
- 매 호출마다 엔진을 생성하고 dispose하므로 오버헤드가 있음
|
|
- 빈번한 호출이 필요한 경우 방법 1(모듈 레벨 엔진)을 고려
|
|
"""
|
|
worker_engine = create_async_engine(
|
|
url=db_settings.MYSQL_URL,
|
|
poolclass=NullPool,
|
|
connect_args={
|
|
"connect_timeout": 3,
|
|
"charset": "utf8mb4",
|
|
},
|
|
)
|
|
session_factory = async_sessionmaker(
|
|
bind=worker_engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False,
|
|
autoflush=False,
|
|
)
|
|
|
|
async with session_factory() as session:
|
|
try:
|
|
yield session
|
|
except Exception as e:
|
|
await session.rollback()
|
|
print(f"Worker session rollback due to: {e}")
|
|
raise e
|
|
finally:
|
|
await session.close()
|
|
|
|
await worker_engine.dispose()
|