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()