File size: 1,813 Bytes
0ba7816 4e91d62 0ba7816 4e91d62 0ba7816 4e91d62 0ba7816 4e91d62 0ba7816 4e91d62 0ba7816 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
"""
Simple database session management for SQLite.
"""
import os
from contextlib import asynccontextmanager
from typing import AsyncGenerator, Optional
from sqlalchemy.ext.asyncio import AsyncSession, AsyncEngine, create_async_engine, async_sessionmaker
from dotenv import load_dotenv
load_dotenv()
# Get DATABASE_URL from environment
DATABASE_URL = os.getenv("DATABASE_URL", "sqlite+aiosqlite:///./cidadao_ai.db")
# Lazy initialization
_engine: Optional[AsyncEngine] = None
_session_factory: Optional[async_sessionmaker] = None
def _get_engine() -> AsyncEngine:
"""Get or create the async engine (lazy initialization)."""
global _engine
if _engine is None:
_engine = create_async_engine(
DATABASE_URL,
echo=False,
pool_pre_ping=True,
)
return _engine
def _get_session_factory() -> async_sessionmaker:
"""Get or create the session factory (lazy initialization)."""
global _session_factory
if _session_factory is None:
_session_factory = async_sessionmaker(
_get_engine(),
class_=AsyncSession,
expire_on_commit=False,
)
return _session_factory
@asynccontextmanager
async def get_db_session() -> AsyncGenerator[AsyncSession, None]:
"""Get database session context manager."""
factory = _get_session_factory()
async with factory() as session:
try:
yield session
await session.commit()
except Exception:
await session.rollback()
raise
finally:
await session.close()
# Alias for FastAPI dependency
async def get_db() -> AsyncGenerator[AsyncSession, None]:
"""FastAPI dependency for database session."""
async with get_db_session() as session:
yield session
|