Spaces:
Running on Zero
Running on Zero
File size: 2,639 Bytes
7982d01 b37ebd1 7982d01 b37ebd1 120bd9a b37ebd1 7982d01 b37ebd1 120bd9a 7982d01 4566c49 b37ebd1 7982d01 b37ebd1 7982d01 b37ebd1 120bd9a 7982d01 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | 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)
@event.listens_for(Session, "after_flush")
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)
|