Spaces:
Running on Zero
Running on Zero
| import asyncio | |
| import logging | |
| from sqlalchemy import event | |
| from sqlalchemy.engine import make_url | |
| from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine | |
| from sqlalchemy.orm import Session | |
| from app.config import get_settings | |
| from app.db.base import Base | |
| from app.db.migrate import run_migrations | |
| log = logging.getLogger(__name__) | |
| def _build_engine(): | |
| url = make_url(get_settings().database_url) | |
| kwargs: dict = {} | |
| connect_args: dict = {} | |
| if url.drivername.startswith("postgresql"): | |
| # Supabase/Neon behind a transaction pooler: | |
| # - drop ?sslmode=..., force asyncpg, require SSL | |
| # - statement_cache_size=0 (pgbouncer transaction mode has no session prepared statements) | |
| # - small client-side pool + pre_ping so stale connections never reach a handler | |
| # - prepared_statement_cache_size=0: SQLAlchemy keeps its OWN prepared-statement cache, which is | |
| # unsafe when the pooler hands each transaction a different backend connection | |
| url = url.set(drivername="postgresql+asyncpg", query={"prepared_statement_cache_size": "0"}) | |
| connect_args = {"ssl": "require", "statement_cache_size": 0} | |
| kwargs.update(pool_size=5, max_overflow=5, pool_pre_ping=True, pool_recycle=1800) | |
| return create_async_engine(url, connect_args=connect_args, **kwargs) | |
| engine = _build_engine() | |
| SessionFactory = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession) | |
| def _mark_written(session, flush_context) -> None: | |
| """Remember that this session wrote something, so the middleware knows a COMMIT is needed | |
| (read-only updates then skip the COMMIT round-trip entirely).""" | |
| session.info["wrote"] = True | |
| async def init_db() -> None: | |
| import app.db.models # noqa: F401 (register tables) | |
| async with engine.begin() as conn: | |
| await conn.run_sync(Base.metadata.create_all) | |
| await run_migrations(engine) | |
| async def warm_up(n: int = 3) -> None: | |
| """Open a few pooled connections now so the first real users don't pay the handshake.""" | |
| async def one() -> None: | |
| async with engine.connect() as conn: | |
| await conn.exec_driver_sql("SELECT 1") | |
| await asyncio.gather(*[one() for _ in range(n)], return_exceptions=True) | |
| async def keepalive(interval: int = 45) -> None: | |
| while True: | |
| await asyncio.sleep(interval) | |
| try: | |
| async with engine.connect() as conn: | |
| await conn.exec_driver_sql("SELECT 1") | |
| except Exception as exc: | |
| log.warning("DB keepalive failed: %r", exc) | |