database/session.py

"""Async database session factory."""

from __future__ import annotations

import os

from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.pool import NullPool

from database.models import Base

DATABASE_URL = os.getenv(
    "DATABASE_URL",
    "postgresql+asyncpg://sentinel:sentinel@localhost:5432/sentinelai",
)

_connect_args: dict = {}
if "postgresql" in DATABASE_URL and "asyncpg" in DATABASE_URL:
    _connect_args["timeout"] = float(os.getenv("DB_CONNECT_TIMEOUT_SEC", "10"))

engine = create_async_engine(
    DATABASE_URL,
    echo=False,
    poolclass=NullPool,
    connect_args=_connect_args or {},
)

async_session_factory = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)


async def init_db() -> None:
    if os.getenv("SKIP_DB", "").lower() in {"1", "true", "yes"}:
        return
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)