""" 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 # Load environment variables from .env file from dotenv import load_dotenv # Get the backend directory path backend_dir = Path(__file__).resolve().parent.parent env_path = backend_dir / '.env' # Load .env file if it exists if env_path.exists(): load_dotenv(dotenv_path=env_path) # Get database URL from environment DATABASE_URL = os.getenv("DATABASE_URL") if not DATABASE_URL: raise ValueError("DATABASE_URL environment variable is not set") # Create async engine for Neon PostgreSQL # echo=True enables SQL query logging (set to False in production) engine = create_async_engine( DATABASE_URL, echo=False, future=True, pool_size=5, max_overflow=10, pool_pre_ping=True, pool_recycle=300, ) # Create async session factory 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 # Import all models to register them with SQLModel from app.models import task, user, conversation, message, dashboard, audit_log async with engine.begin() as conn: # Create all tables defined in SQLModel models await conn.run_sync(SQLModel.metadata.create_all)