| """
|
| Task: T007
|
| Spec: 002-authentication/data-model.md - Database Connection
|
|
|
| SQLModel async session management for Neon PostgreSQL database.
|
| """
|
|
|
| import os
|
| from pathlib import Path
|
| from sqlmodel import create_engine, Session
|
| from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
|
| from typing import AsyncGenerator
|
|
|
|
|
| from dotenv import load_dotenv
|
|
|
|
|
| backend_dir = Path(__file__).resolve().parent.parent
|
| env_path = backend_dir / '.env'
|
|
|
|
|
| if env_path.exists():
|
| load_dotenv(dotenv_path=env_path)
|
|
|
|
|
| DATABASE_URL = os.getenv("DATABASE_URL")
|
|
|
| if not DATABASE_URL:
|
| raise ValueError("DATABASE_URL environment variable is not set")
|
|
|
|
|
|
|
| engine = create_async_engine(
|
| DATABASE_URL,
|
| echo=False,
|
| future=True,
|
| pool_size=5,
|
| max_overflow=10,
|
| pool_pre_ping=True,
|
| pool_recycle=300,
|
| )
|
|
|
|
|
| async_session_maker = async_sessionmaker(
|
| engine,
|
| class_=AsyncSession,
|
| expire_on_commit=False
|
| )
|
|
|
|
|
| async def get_session() -> AsyncGenerator[AsyncSession, None]:
|
| """
|
| Dependency to get async database session.
|
|
|
| Usage in FastAPI endpoints:
|
| @router.get("/endpoint")
|
| async def endpoint(session: AsyncSession = Depends(get_session)):
|
| # Use session here
|
| pass
|
|
|
| Yields:
|
| AsyncSession: SQLAlchemy async session
|
| """
|
| async with async_session_maker() as session:
|
| yield session
|
|
|
|
|
| async def init_db():
|
| """
|
| Initialize database tables.
|
|
|
| Note: In production, use Alembic migrations instead.
|
| This is useful for testing or initial setup.
|
| """
|
| from sqlmodel import SQLModel
|
|
|
| from app.models import task, user, conversation, message, dashboard, audit_log
|
|
|
| async with engine.begin() as conn:
|
|
|
| await conn.run_sync(SQLModel.metadata.create_all)
|
|
|
