teecup/app/db.py

"""
Databaselag for TeeCup.

Det kritiske her er hvordan organisasjonskonteksten (`app.current_org`) settes:

  * Den settes med set_config(..., is_local => true), dvs. SET LOCAL — verdien
    gjelder KUN inne i den aktive transaksjonen og nullstilles automatisk når
    transaksjonen avsluttes.
  * Uten dette (en vanlig SET på en pool-tilkobling) ville konteksten blitt
    liggende på tilkoblingen når den returneres til poolen, og neste forespørsel
    — potensielt en ANNEN organisasjon — ville arvet den. Det ville brutt
    isolasjonen ved gjenbruk av tilkoblinger.

Hver forespørsel som rører organisasjonsdata kjører altså i én transaksjon med
org-konteksten satt lokalt.
"""

from contextlib import asynccontextmanager
from typing import AsyncIterator

import asyncpg

from .config import settings

_pool: asyncpg.Pool | None = None


async def init_pool() -> None:
    global _pool
    _pool = await asyncpg.create_pool(
        host=settings.DB_HOST,
        port=settings.DB_PORT,
        user=settings.DB_USER,
        password=settings.DB_PASS,
        database=settings.DB_NAME,
        min_size=settings.DB_POOL_MIN,
        max_size=settings.DB_POOL_MAX,
    )


async def close_pool() -> None:
    global _pool
    if _pool is not None:
        await _pool.close()
        _pool = None


def _require_pool() -> asyncpg.Pool:
    if _pool is None:
        raise RuntimeError("DB-poolen er ikke initialisert (kall init_pool ved oppstart).")
    return _pool


@asynccontextmanager
async def plain_connection() -> AsyncIterator[asyncpg.Connection]:
    """Tilkobling UTEN org-kontekst.

    Brukes kun til auth-laget (f.eks. oppslag i organization_membership, som
    ikke er organisasjonsavgrenset). Rør aldri organisasjonsdata her.
    """
    pool = _require_pool()
    async with pool.acquire() as conn:
        yield conn


@asynccontextmanager
async def org_connection(organization_id: str) -> AsyncIterator[asyncpg.Connection]:
    """Tilkobling med org-kontekst satt lokalt for én transaksjon.

    All lesing/skriving av organisasjonsdata skal gå gjennom denne. RLS filtrerer
    da automatisk på organization_id = app.current_org.

    VIKTIG: `organization_id` MÅ på forhånd være verifisert mot brukerens
    medlemskap (se auth.get_authorized_org). Denne funksjonen stoler på at
    verdien allerede er autorisert.
    """
    pool = _require_pool()
    async with pool.acquire() as conn:
        async with conn.transaction():
            # SET LOCAL via set_config(..., true): nullstilles ved transaksjonsslutt.
            await conn.execute(
                "SELECT set_config('app.current_org', $1, true)",
                str(organization_id),
            )
            yield conn