Files
rose-ash/shared/db/session.py
giles 961067841e
All checks were successful
Build and Deploy / build-and-deploy (push) Successful in 3m14s
Tier 0 scalability: PgBouncer, Redis split, DB split, workers
T0.1: Separate redis-auth service (64mb, noeviction) for auth state
T0.2: Bump data Redis from 256mb to 1gb
T0.3: Per-app DATABASE_URL via PgBouncer to per-domain databases
T0.4: PgBouncer service (transaction mode, pool=20, max_conn=300);
      session.py pools reduced to 3+5 with timeout and recycle
T0.5: Hypercorn --workers 2 --keep-alive 75 on all 6 entrypoints

Deploy requires running split-databases.sh first to create per-domain
databases from the existing appdb.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-25 10:10:18 +00:00

153 lines
4.2 KiB
Python

from __future__ import annotations
import os
from contextlib import asynccontextmanager
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
from quart import Quart, g
DATABASE_URL = (
os.getenv("DATABASE_URL_ASYNC")
or os.getenv("DATABASE_URL")
or "postgresql+asyncpg://localhost/blog"
)
_engine = create_async_engine(
DATABASE_URL,
future=True,
echo=False,
pool_pre_ping=True,
pool_size=int(os.getenv("DB_POOL_SIZE", "3")),
max_overflow=int(os.getenv("DB_MAX_OVERFLOW", "5")),
pool_timeout=10,
pool_recycle=1800,
)
_Session = async_sessionmaker(
bind=_engine,
class_=AsyncSession,
expire_on_commit=False,
)
@asynccontextmanager
async def get_session():
"""Always create a fresh AsyncSession for this block."""
sess = _Session()
try:
yield sess
finally:
await sess.close()
# ---------------------------------------------------------------------------
# Cross-domain sessions — account + federation
#
# Initially DATABASE_URL_ACCOUNT / DATABASE_URL_FEDERATION point to the same
# DB as DATABASE_URL (zero behaviour change). When per-domain DBs are ready,
# switch the env vars to the new connection strings.
# ---------------------------------------------------------------------------
DATABASE_URL_ACCOUNT = (
os.getenv("DATABASE_URL_ACCOUNT") or DATABASE_URL
)
DATABASE_URL_FEDERATION = (
os.getenv("DATABASE_URL_FEDERATION") or DATABASE_URL
)
# Engines are created lazily — only allocate a pool if the URL differs
_account_engine = (
_engine if DATABASE_URL_ACCOUNT == DATABASE_URL
else create_async_engine(
DATABASE_URL_ACCOUNT,
future=True, echo=False, pool_pre_ping=True,
pool_size=2, max_overflow=3,
pool_timeout=10, pool_recycle=1800,
)
)
_AccountSession = async_sessionmaker(
bind=_account_engine,
class_=AsyncSession,
expire_on_commit=False,
)
_federation_engine = (
_engine if DATABASE_URL_FEDERATION == DATABASE_URL
else create_async_engine(
DATABASE_URL_FEDERATION,
future=True, echo=False, pool_pre_ping=True,
pool_size=2, max_overflow=3,
pool_timeout=10, pool_recycle=1800,
)
)
_FederationSession = async_sessionmaker(
bind=_federation_engine,
class_=AsyncSession,
expire_on_commit=False,
)
@asynccontextmanager
async def get_account_session():
"""Session targeting the account database (users, grants, oauth codes)."""
sess = _AccountSession()
try:
yield sess
finally:
await sess.close()
@asynccontextmanager
async def get_federation_session():
"""Session targeting the federation database (ap_activities, etc.)."""
sess = _FederationSession()
try:
yield sess
finally:
await sess.close()
def register_db(app: Quart):
@app.before_request
async def open_session():
g.s = _Session()
g.tx = await g.s.begin()
g.had_error = False
@app.after_request
async def maybe_commit(response):
# Runs BEFORE bytes are sent.
if not g.had_error and 200 <= response.status_code < 400:
try:
if hasattr(g, "tx"):
await g.tx.commit()
except Exception as e:
print(f'commit failed {e}')
if hasattr(g, "tx"):
await g.tx.rollback()
from quart import make_response
return await make_response("Commit failed", 500)
return response
@app.teardown_request
async def finish(exc):
try:
# If an exception occurred OR we didn't commit (still in txn), roll back.
if hasattr(g, "s"):
if exc is not None or g.s.in_transaction():
if hasattr(g, "tx") and g.tx.is_active:
try:
await g.tx.rollback()
except Exception:
pass
finally:
if hasattr(g, "s"):
try:
await g.s.close()
except Exception:
pass
@app.errorhandler(Exception)
async def mark_error(e):
g.had_error = True
raise