jpmschweitzerandClaude Opus 4.5
2523db4da7
- SQLAlchemy async database layer (SQLite dev, PostgreSQL prod) - Conversation and Message models with UUID primary keys - Token counting utilities using litellm - Context summarization at 80% token threshold - REST API endpoints for multi-turn conversations - 19 conversation tests, 6 token tests (176 total passing) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
124 lines
3.4 KiB
Python
124 lines
3.4 KiB
Python
"""
|
|
Async SQLAlchemy database management.
|
|
|
|
Pattern from core-api: singleton Database class with async session factory.
|
|
"""
|
|
from collections.abc import AsyncGenerator
|
|
from functools import lru_cache
|
|
|
|
from sqlalchemy.ext.asyncio import (
|
|
AsyncEngine,
|
|
AsyncSession,
|
|
async_sessionmaker,
|
|
create_async_engine,
|
|
)
|
|
|
|
from src.shared.config import get_settings
|
|
from src.shared.logging import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class Database:
|
|
"""
|
|
Async database connection manager.
|
|
|
|
Manages SQLAlchemy async engine and session factory.
|
|
"""
|
|
|
|
def __init__(self, url: str):
|
|
"""
|
|
Initialize database with connection URL.
|
|
|
|
Args:
|
|
url: SQLAlchemy async connection URL
|
|
e.g., "sqlite+aiosqlite:///./webber.db"
|
|
or "postgresql+asyncpg://user:pass@host/db"
|
|
"""
|
|
self._url = url
|
|
self._engine: AsyncEngine | None = None
|
|
self._session_factory: async_sessionmaker[AsyncSession] | None = None
|
|
|
|
@property
|
|
def engine(self) -> AsyncEngine:
|
|
"""Get or create the async engine."""
|
|
if self._engine is None:
|
|
self._engine = create_async_engine(
|
|
self._url,
|
|
echo=get_settings().debug,
|
|
pool_pre_ping=True,
|
|
)
|
|
return self._engine
|
|
|
|
@property
|
|
def session_factory(self) -> async_sessionmaker[AsyncSession]:
|
|
"""Get or create the session factory."""
|
|
if self._session_factory is None:
|
|
self._session_factory = async_sessionmaker(
|
|
bind=self.engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False,
|
|
autoflush=False,
|
|
)
|
|
return self._session_factory
|
|
|
|
async def create_tables(self) -> None:
|
|
"""Create all tables (for development)."""
|
|
from src.db.models import Base
|
|
|
|
async with self.engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
logger.info("Database tables created")
|
|
|
|
async def close(self) -> None:
|
|
"""Close the database connection."""
|
|
if self._engine:
|
|
await self._engine.dispose()
|
|
self._engine = None
|
|
self._session_factory = None
|
|
logger.info("Database connection closed")
|
|
|
|
|
|
# Singleton instance
|
|
_database: Database | None = None
|
|
_tables_created: bool = False
|
|
|
|
|
|
@lru_cache
|
|
def get_database() -> Database:
|
|
"""Get the database singleton."""
|
|
global _database
|
|
if _database is None:
|
|
settings = get_settings()
|
|
_database = Database(settings.database_url)
|
|
return _database
|
|
|
|
|
|
async def _ensure_tables() -> None:
|
|
"""Ensure database tables exist (lazy initialization)."""
|
|
global _tables_created
|
|
if not _tables_created:
|
|
database = get_database()
|
|
await database.create_tables()
|
|
_tables_created = True
|
|
|
|
|
|
async def get_session() -> AsyncGenerator[AsyncSession, None]:
|
|
"""
|
|
Dependency for getting async database sessions.
|
|
|
|
Usage:
|
|
@router.get("/")
|
|
async def endpoint(session: AsyncSession = Depends(get_session)):
|
|
...
|
|
"""
|
|
await _ensure_tables()
|
|
database = get_database()
|
|
async with database.session_factory() as session:
|
|
try:
|
|
yield session
|
|
await session.commit()
|
|
except Exception:
|
|
await session.rollback()
|
|
raise
|