app/tenancy.py
Rama main · réplica de lectura del repositorio original.
"""Automatic per-portfolio filtering of the tables holding personal data.
Every `select()` on a session carrying a portfolio id gets a
`WHERE portfolio_id = :pid` on the models in `SCOPED_MODELS` — top-level
queries, eager loads and lazy loads alike.
It fails closed: a select touching a scoped model with no portfolio id on the
session raises. `unscoped_session()` is the only opt-out, and is greppable.
Writes are not covered — SQLAlchemy applies loader criteria to selects only —
so inserts set `portfolio_id` via `owned()`, with a before_flush net in
`install()`.
This separates users *within the application*, not from anyone holding the
database: see decisions.md #63 (Trust boundary).
"""
from __future__ import annotations
from contextlib import contextmanager
from typing import Iterator
from sqlalchemy import event, select
from sqlalchemy import false as sql_false
from sqlalchemy.orm import Session, with_loader_criteria
from .models import (
MEMBER_ROLES,
CorporateAction,
Dividend,
HoldingPref,
InvestmentPlan,
PlannedPurchase,
PortfolioMember,
SavedChart,
Trade,
)
# Every model with a portfolio_id column. Adding a personal-data table? Add it
# here — that is the whole registration step.
SCOPED_MODELS = (
Trade,
Dividend,
CorporateAction,
PlannedPurchase,
HoldingPref,
InvestmentPlan,
)
# Owned by a person rather than a portfolio (their charts follow them between
# portfolios). Filtered the same way, on user_id — being in this tuple is what
# keeps them inside the fail-closed guarantee rather than relying on every
# query remembering a where clause.
USER_SCOPED_MODELS = (SavedChart,)
_PORTFOLIO_KEY = "portfolio_id"
_USER_KEY = "portfolio_user_id"
_BYPASS_KEY = "portfolio_unscoped"
class TenancyError(RuntimeError):
"""Raised when personal data is queried with no portfolio context."""
_UNBOUND_MESSAGE = (
"query touches portfolio data with no portfolio on the session — use "
"portfolio_session(...) or, for market-data maintenance, unscoped_session(...)"
)
def install(session_factory) -> None:
"""Attach the filter to a sessionmaker. Call once at startup."""
@event.listens_for(session_factory, "before_flush")
def _stamp_owner(session, _ctx, _instances) -> None:
"""New personal rows inherit the session's portfolio, and record who
entered them.
`owned()` says it explicitly at the call sites that matter; this is the
safety net for bulk paths (the workbook importer creates rows in eight
places) and for anything added later. Only fills blanks — it never
reassigns an existing owner.
"""
portfolio_id = session.info.get(_PORTFOLIO_KEY)
user_id = session.info.get(_USER_KEY)
if portfolio_id is None:
return
for obj in session.new:
if not isinstance(obj, SCOPED_MODELS):
continue
if obj.portfolio_id is None:
obj.portfolio_id = portfolio_id
if user_id is not None and getattr(obj, "user_id", "absent") is None:
obj.user_id = user_id
@event.listens_for(session_factory, "do_orm_execute")
def _apply_portfolio_filter(state) -> None:
if not state.is_select or state.is_column_load:
return
session = state.session
if session.info.get(_BYPASS_KEY):
return
portfolio_id = session.info.get(_PORTFOLIO_KEY)
# A relationship load inherits the criteria from the query that loaded
# its parent (propagate_to_loaders), so re-applying them here would
# double up. But an UNBOUND session has no criteria to inherit, and a
# lazy `instrument.trades` would then hand back every portfolio's rows
# — so the refusal below still has to run before we return.
if state.is_relationship_load:
if portfolio_id is None and _touches_scoped(state):
raise TenancyError(_UNBOUND_MESSAGE)
return
if portfolio_id is None:
# No portfolio bound: fine for market data (the price feed runs this
# way), a bug for anything personal.
if _touches_scoped(state):
raise TenancyError(_UNBOUND_MESSAGE)
return
# Applied to EVERY statement, not just ones naming a scoped model:
# `select(Instrument)` carries the criteria forward to a later
# `instrument.trades` lazy load (propagate_to_loaders). Skipping it here
# because the statement "doesn't touch trades" is exactly how that lazy
# load would come back unfiltered.
for model in SCOPED_MODELS:
state.statement = state.statement.options(
with_loader_criteria(
model, model.portfolio_id == portfolio_id, include_aliases=True
)
)
# User-scoped models are filtered whether or not a user is bound. With
# no user, the criterion matches nothing rather than everything: a
# session carrying a portfolio but no user (an API key whose creator
# was deleted, nulling `created_by`) would otherwise read EVERY user's
# saved charts.
user_id = session.info.get(_USER_KEY)
for model in USER_SCOPED_MODELS:
criterion = (
model.user_id == user_id if user_id is not None else sql_false()
)
state.statement = state.statement.options(
with_loader_criteria(model, criterion, include_aliases=True)
)
def _touches_scoped(state) -> bool:
"""True when the statement selects, or eager-loads, a scoped model.
Cheap over-approximation: match on the mapper of every entity in the
statement. Instrument-only queries (price feed, catalogue) skip the check so
they keep working outside a request.
"""
scoped = set(SCOPED_MODELS) | set(USER_SCOPED_MODELS)
for desc in state.statement.column_descriptions:
if desc.get("entity") in scoped:
return True
# A scoped table reached through FROM or JOIN rather than named as the
# selected entity: `select(func.count()).select_from(Trade)` selects a
# number, not a Trade, and `select(Instrument).join(Trade)` selects an
# instrument — both read personal rows and neither shows up above.
# Compared by table, because that is what a JOIN carries.
scoped_tables = {model.__table__ for model in scoped}
for element in getattr(state.statement, "get_final_froms", lambda: ())():
if _mentions(element, scoped_tables):
return True
# Eager loads name their target in the statement's _with_options; simplest
# reliable signal is the ORM path they walk, every step of it, so a scoped
# model in the middle of a nested load counts too.
for opt in getattr(state.statement, "_with_options", ()):
for attr in getattr(opt, "path", ()) or ():
if getattr(getattr(attr, "entity", None), "class_", None) in scoped:
return True
return False
def _mentions(element, tables: set) -> bool:
"""Whether a FROM element is, or joins, one of these tables."""
if element in tables:
return True
left, right = getattr(element, "left", None), getattr(element, "right", None)
if left is not None and _mentions(left, tables):
return True
if right is not None and _mentions(right, tables):
return True
# A subquery keeps its own SELECT, which may reach a scoped table too.
inner = getattr(getattr(element, "element", None), "get_final_froms", None)
if inner is not None:
return any(_mentions(f, tables) for f in inner())
return False
def bind(session: Session, portfolio_id: int, user_id: int | None = None) -> Session:
"""Confine an already-open session to one portfolio's rows.
`user_id` is stamped onto new rows as "recorded by" — it plays no part in
what the session can read.
"""
session.info[_PORTFOLIO_KEY] = portfolio_id
session.info[_USER_KEY] = user_id
return session
def allow_unscoped(session: Session) -> Session:
"""Drop the filter on an open session. Two callers only: the first-run
claim of pre-auth rows (which have no owner yet, so no filter could find
them) and the price feed's cross-portfolio FX repair."""
session.info[_BYPASS_KEY] = True
return session
@contextmanager
def as_portfolio(
session: Session,
portfolio_id: int,
*,
user_id: int | None,
roles: tuple[str, ...] = MEMBER_ROLES,
) -> Iterator[Session]:
"""Read another of THIS USER's portfolios for the length of the block.
The one operation that spans two portfolios is moving a holding between
them, and it must validate the destination's timeline before writing there.
`unscoped_session` is the wrong tool for it: that drops the filter for
every portfolio at once and says never to serve a request with it.
The membership check is HERE, not in the caller. A primitive that rebinds
the filter on request is a cross-tenant read waiting for one careless call
site — so it refuses unless `user_id` holds a role in `roles`, and callers
that are about to WRITE pass `roles=WRITE_ROLES` rather than checking
separately and hoping the two agree.
Reads inside the block must use `select()`. `session.get()` can answer from
the identity map without going near the filter, so it would hand back a row
belonging to whichever portfolio loaded it first.
"""
if user_id is None:
raise TenancyError("cannot reach another portfolio with no user")
role = session.scalar(
select(PortfolioMember.role).where(
PortfolioMember.portfolio_id == portfolio_id,
PortfolioMember.user_id == user_id,
)
)
if role is None or role not in roles:
# One message for "not a member" and "not allowed to write", because
# telling them apart would confirm the portfolio exists.
raise TenancyError(
f"user {user_id} may not reach portfolio {portfolio_id} as {roles}"
)
was = session.info.get(_PORTFOLIO_KEY)
session.info[_PORTFOLIO_KEY] = portfolio_id
try:
yield session
finally:
# Restored even when the block raises: a failed validation must not
# leave the rest of the request reading the wrong portfolio.
session.info[_PORTFOLIO_KEY] = was
@contextmanager
def portfolio_session(
session_factory, portfolio_id: int, user_id: int | None = None
) -> Iterator[Session]:
"""A session whose reads are confined to one portfolio. Commits on success."""
session = session_factory()
session.info[_PORTFOLIO_KEY] = portfolio_id
session.info[_USER_KEY] = user_id
try:
yield session
session.commit()
except Exception:
session.rollback()
raise
finally:
session.close()
@contextmanager
def unscoped_session(session_factory) -> Iterator[Session]:
"""Every portfolio's rows — market-data maintenance only (the price feed's
trade-FX backfill). Never use this to serve a request."""
session = session_factory()
session.info[_BYPASS_KEY] = True
try:
yield session
session.commit()
except Exception:
session.rollback()
raise
finally:
session.close()
def current_portfolio_id(session: Session) -> int | None:
return session.info.get(_PORTFOLIO_KEY)
def current_user_id(session: Session) -> int | None:
return session.info.get(_USER_KEY)
def owned(session: Session, obj):
"""Stamp a new row with the session's portfolio (and recorder) before add."""
from .models import SavedChart as _SavedChart
if isinstance(obj, _SavedChart):
user_id = current_user_id(session)
if user_id is None:
raise TenancyError("cannot create a chart with no user context")
obj.user_id = user_id
obj.portfolio_id = current_portfolio_id(session)
return obj
portfolio_id = current_portfolio_id(session)
if portfolio_id is None:
raise TenancyError(
f"cannot create {type(obj).__name__} with no portfolio context"
)
obj.portfolio_id = portfolio_id
user_id = current_user_id(session)
if user_id is not None and hasattr(obj, "user_id"):
obj.user_id = user_id
return obj