osint-dashboard/app/database.py

import os

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

from config import DATABASE_URL

# Pooled connections are bound to the event loop they were created on. Long-lived
# servers (uvicorn, the ingester) want a pool; short-lived / test processes that
# open a fresh event loop per unit (asyncio.run, pytest) must not reuse pooled
# connections across loops, so allow a NullPool (new connection per session).
_NULL_POOL = os.getenv("DB_NULL_POOL", "").lower() in ("1", "true", "yes")

engine = create_async_engine(
    DATABASE_URL,
    echo=False,
    **({"poolclass": NullPool} if _NULL_POOL else
       {"pool_size": 5, "max_overflow": 10, "pool_recycle": 300}),
)
async_session = async_sessionmaker(
    engine, class_=AsyncSession, expire_on_commit=False
)
metadata = MetaData()


async def init_extensions():
    """Initialize PostGIS and TimescaleDB extensions on first connection."""
    async with engine.connect() as conn:
        await conn.execute(text("CREATE EXTENSION IF NOT EXISTS postgis"))
        await conn.execute(text("CREATE EXTENSION IF NOT EXISTS timescaledb"))
        await conn.commit()