| import os
|
| import sys
|
| from contextlib import asynccontextmanager
|
| import asyncio
|
| import time
|
| from sqlalchemy.exc import DBAPIError
|
| from datetime import datetime
|
| from typing import AsyncIterator
|
|
|
| from sqlalchemy import func
|
| from sqlalchemy.engine import URL
|
| from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
| from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
|
|
| DATABASE_URL = os.getenv("DATABASE_URL")
|
|
|
| if not DATABASE_URL:
|
| raise ValueError("DATABASE_URL environment variable is not set.")
|
| else:
|
| DATABASE_URL = str(DATABASE_URL)
|
|
|
|
|
| POOL_SIZE = int(os.getenv("DB_POOL_SIZE", 5))
|
| MAX_OVERFLOW = int(os.getenv("DB_MAX_OVERFLOW", 10))
|
| POOL_TIMEOUT = int(os.getenv("DB_POOL_TIMEOUT", 30))
|
|
|
| connect_args = (
|
| {"check_same_thread": False} if DATABASE_URL and "sqlite" in DATABASE_URL else {}
|
| )
|
|
|
| engine = create_async_engine(
|
| DATABASE_URL,
|
| connect_args=connect_args,
|
| pool_pre_ping=True,
|
| pool_size=POOL_SIZE,
|
| max_overflow=MAX_OVERFLOW,
|
| pool_timeout=POOL_TIMEOUT,
|
| )
|
|
|
| session_pool: async_sessionmaker[AsyncSession] = async_sessionmaker(
|
| bind=engine, autocommit=False, autoflush=False, expire_on_commit=False
|
| )
|
|
|
|
|
| class Base(DeclarativeBase):
|
| """Base class for declarative models."""
|
|
|
| pass
|
|
|
|
|
| class TimestampMixin:
|
| """Mixin to add created_at and updated_at timestamps to models."""
|
|
|
| created_at: Mapped[datetime] = mapped_column(
|
| server_default=func.now(), nullable=False
|
| )
|
| updated_at: Mapped[datetime] = mapped_column(
|
| server_default=func.now(), onupdate=func.now(), nullable=False
|
| )
|
|
|
|
|
| @asynccontextmanager
|
| async def get_session() -> AsyncIterator[AsyncSession]:
|
| """
|
| Context manager to get a database session.
|
| Provides a session from the connection pool. The session automatically
|
| handles transaction lifecycle - commits on success, rolls back on exceptions.
|
| The session is properly closed after use to return it to the pool.
|
|
|
| Usage:
|
| ```
|
| async with get_session() as session:
|
| result = await session.execute(stmt)
|
| await session.commit() # Explicit commit if needed
|
| ```
|
| """
|
|
|
| retries = int(os.getenv("DB_SESSION_RETRIES", 2))
|
| delay = float(os.getenv("DB_SESSION_RETRY_DELAY", 0.2))
|
|
|
| last_exc = None
|
| for attempt in range(retries + 1):
|
| try:
|
| async with session_pool() as session:
|
| yield session
|
| return
|
| except DBAPIError as e:
|
| last_exc = e
|
|
|
| if attempt < retries:
|
| await asyncio.sleep(delay)
|
| continue
|
| raise
|
|
|
