diff options
Diffstat (limited to 'app/utils/db.py')
| -rw-r--r-- | app/utils/db.py | 49 |
1 files changed, 32 insertions, 17 deletions
diff --git a/app/utils/db.py b/app/utils/db.py index 5531998..d3c7318 100644 --- a/app/utils/db.py +++ b/app/utils/db.py @@ -1,7 +1,6 @@ import os -from typing import AsyncGenerator +from typing import AsyncGenerator, Optional -from dotenv import load_dotenv from sqlalchemy.ext.asyncio import ( AsyncEngine, AsyncSession, @@ -10,29 +9,45 @@ from sqlalchemy.ext.asyncio import ( ) from sqlalchemy.orm import DeclarativeBase -load_dotenv() +engine: Optional[AsyncEngine] = None +async_session: Optional[async_sessionmaker[AsyncSession]] = None -DATABASE_URL = os.getenv("DATABASE_URL") -if not DATABASE_URL: - raise ValueError("DATABASE_URL not found in .env") +class Base(DeclarativeBase): + pass -engine: AsyncEngine = create_async_engine( - DATABASE_URL, - echo=True, - pool_pre_ping=True, -) -async_session = async_sessionmaker( - bind=engine, - expire_on_commit=False, -) +def init_db_engine() -> None: + global engine, async_session + database_url = os.getenv("DATABASE_URL") -class Base(DeclarativeBase): - pass + if not database_url: + raise RuntimeError("DATABASE_URL not found in environment") + + engine = create_async_engine( + database_url, + echo=True, + pool_pre_ping=True, + ) + + async_session = async_sessionmaker( + bind=engine, + expire_on_commit=False, + ) + + +def get_engine() -> AsyncEngine: + if engine is None: + raise RuntimeError( + "DB engine not initialized. Call init_db_engine() first." + ) + return engine async def get_async_session() -> AsyncGenerator[AsyncSession, None]: + if async_session is None: + raise RuntimeError("DB not initialized. Call init_db_engine() first.") + async with async_session() as session: yield session |
