64 lines
2.1 KiB
Python
64 lines
2.1 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""비동기 DB 세션 + Base + 공통 타임스탬프 믹스인.
|
|
|
|
castad 는 공통 base 가 없어 컬럼을 매번 반복하지만, 여기선 TimestampMixin 으로 정리.
|
|
"""
|
|
from collections.abc import AsyncGenerator
|
|
from datetime import datetime
|
|
|
|
from sqlalchemy import DateTime, func
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
|
|
|
from .config import settings
|
|
|
|
engine = create_async_engine(settings.DATABASE_URL, echo=False, future=True)
|
|
SessionLocal = async_sessionmaker(bind=engine, class_=AsyncSession, expire_on_commit=False, autoflush=False)
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
pass
|
|
|
|
|
|
class TimestampMixin:
|
|
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
|
|
|
|
|
|
async def get_session() -> AsyncGenerator[AsyncSession, None]:
|
|
async with SessionLocal() as session:
|
|
try:
|
|
yield session
|
|
except Exception:
|
|
await session.rollback()
|
|
raise
|
|
|
|
|
|
async def create_tables() -> None:
|
|
from . import models # noqa: F401 (테이블 메타데이터 등록)
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
await _lightweight_migrate(conn)
|
|
|
|
|
|
# create_all 은 기존 테이블에 새 컬럼을 추가하지 않는다. 개발용 SQLite 를 위한 최소 마이그레이션.
|
|
# (운영 MySQL 은 Alembic 등 정식 마이그레이션 사용 권장)
|
|
_ADD_COLUMNS = {
|
|
"user": [
|
|
("profile_image_url", "TEXT"),
|
|
("bio", "VARCHAR(200)"),
|
|
],
|
|
}
|
|
|
|
|
|
async def _lightweight_migrate(conn) -> None:
|
|
from sqlalchemy import text
|
|
dialect = conn.dialect.name
|
|
if dialect != "sqlite":
|
|
return # 운영 DB 는 정식 마이그레이션으로 처리
|
|
for table, cols in _ADD_COLUMNS.items():
|
|
rows = (await conn.exec_driver_sql(f"PRAGMA table_info({table})")).fetchall()
|
|
existing = {r[1] for r in rows}
|
|
for name, ddl in cols:
|
|
if name not in existing:
|
|
await conn.exec_driver_sql(f"ALTER TABLE {table} ADD COLUMN {name} {ddl}")
|