97 lines
3.1 KiB
Python
97 lines
3.1 KiB
Python
|
"""Which workspace the current unit of work belongs to.
|
||
|
|
|
||
|
|
Held in a `ContextVar` rather than passed around, because the thing that needs it
|
||
|
|
is the database session listener in `app.core.rls`, and threading a tenant id
|
||
|
|
through every call site to reach it is exactly the discipline that failed —
|
||
|
|
finding S-5 was one query that forgot.
|
||
|
|
|
||
|
|
Three states, and the difference between them matters:
|
||
|
|
|
||
|
|
- **scoped** — a workspace is set. Queries see only that workspace's rows.
|
||
|
|
- **unscoped** — deliberately crossing workspaces. Used by sign-in (which looks a
|
||
|
|
user up by email before any workspace is known), by the event worker, and by
|
||
|
|
platform superadmin operations. Always explicit, never a default.
|
||
|
|
- **unset** — no context established. Queries see **nothing**.
|
||
|
|
|
||
|
|
That last one is the important choice. Unset could have meant "see everything",
|
||
|
|
which would be forgiving of a code path that forgot to set context — and would
|
||
|
|
also mean any such path silently leaks across workspaces. Seeing nothing makes a
|
||
|
|
missed path fail loudly and safely instead.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import functools
|
||
|
|
import uuid
|
||
|
|
from collections.abc import Callable, Iterator
|
||
|
|
from contextlib import contextmanager
|
||
|
|
from contextvars import ContextVar, Token
|
||
|
|
from typing import Any, TypeVar
|
||
|
|
|
||
|
|
_tenant_id: ContextVar[uuid.UUID | None] = ContextVar("saas_tenant_id", default=None)
|
||
|
|
_bypass: ContextVar[bool] = ContextVar("saas_tenant_bypass", default=False)
|
||
|
|
|
||
|
|
F = TypeVar("F", bound=Callable[..., Any])
|
||
|
|
|
||
|
|
|
||
|
|
def current_tenant_id() -> uuid.UUID | None:
|
||
|
|
return _tenant_id.get()
|
||
|
|
|
||
|
|
|
||
|
|
def is_bypassed() -> bool:
|
||
|
|
return _bypass.get()
|
||
|
|
|
||
|
|
|
||
|
|
@contextmanager
|
||
|
|
def scoped_to(tenant_id: uuid.UUID | None) -> Iterator[None]:
|
||
|
|
"""Run inside one workspace.
|
||
|
|
|
||
|
|
A `None` tenant here means a platform superadmin, who belongs to no workspace
|
||
|
|
— so it bypasses rather than scoping to nothing. Scoping a superadmin to
|
||
|
|
"no workspace" would make every administrative screen empty.
|
||
|
|
"""
|
||
|
|
if tenant_id is None:
|
||
|
|
with unscoped():
|
||
|
|
yield
|
||
|
|
return
|
||
|
|
|
||
|
|
token_tenant = _tenant_id.set(tenant_id)
|
||
|
|
token_bypass = _bypass.set(False)
|
||
|
|
try:
|
||
|
|
yield
|
||
|
|
finally:
|
||
|
|
_tenant_id.reset(token_tenant)
|
||
|
|
_bypass.reset(token_bypass)
|
||
|
|
|
||
|
|
|
||
|
|
@contextmanager
|
||
|
|
def unscoped() -> Iterator[None]:
|
||
|
|
"""Deliberately cross workspace boundaries.
|
||
|
|
|
||
|
|
Every use is a decision worth reading twice. The legitimate ones are narrow:
|
||
|
|
authenticating someone before their workspace is known, background work that
|
||
|
|
spans workspaces, and platform superadmin operations.
|
||
|
|
"""
|
||
|
|
token_tenant = _tenant_id.set(None)
|
||
|
|
token_bypass = _bypass.set(True)
|
||
|
|
try:
|
||
|
|
yield
|
||
|
|
finally:
|
||
|
|
_tenant_id.reset(token_tenant)
|
||
|
|
_bypass.reset(token_bypass)
|
||
|
|
|
||
|
|
|
||
|
|
def system_operation(fn: F) -> F:
|
||
|
|
"""Mark a function as crossing workspaces by nature.
|
||
|
|
|
||
|
|
Reads better than an `unscoped()` block at the top of a body, and makes the
|
||
|
|
property visible at the call site rather than buried in the implementation.
|
||
|
|
"""
|
||
|
|
|
||
|
|
@functools.wraps(fn)
|
||
|
|
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
||
|
|
with unscoped():
|
||
|
|
return fn(*args, **kwargs)
|
||
|
|
|
||
|
|
return wrapper # type: ignore[return-value]
|