mono/shared/db/session.py

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 needs_federation_session() -> bool:
    """True when the federation DB is separate from the app's main DB."""
    return DATABASE_URL_FEDERATION != DATABASE_URL


def create_federation_session() -> AsyncSession:
    """Create an unmanaged session targeting the federation database."""
    return _FederationSession()


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