Steven
0c53aa7062
- backend: FastAPI 应用工厂、统一错误信封、JWT、短信登录、工作区/成员/邀请、幂等中间件 - alembic 0001: workspace/member/invitation/audit_log/app_user + RLS 策略 + 双角色 GRANT - tests: 认证/工作区/租户隔离/幂等 四组用例（尚未在本机执行） - web: React18+Vite 骨架（登录、工作区、邀请落地页） - android: Compose 骨架（登录、工作区列表） - 2.Design: PRODUCT-PLAN / ROADMAP-M0-M1 / NAMING / DEVELOPMENT
95 lines
3.5 KiB
Python
95 lines
3.5 KiB
Python
"""数据库会话与租户上下文。
|
|
|
|
安全模型（务必理解，否则会写出越权 bug）：
|
|
|
|
1. 应用进程用 **dadao_app** 角色连库 —— 非 owner、非 superuser、非 bypassrls。
|
|
owner/superuser 会**绕过 RLS**，所以应用绝不使用迁移角色。
|
|
2. 每个请求在**同一个事务**里先 `set_config('app.ws_id', ..., true)`
|
|
（true = 事务级，事务结束自动失效，不会污染连接池里的下一条请求）。
|
|
3. RLS 策略拿这个变量做过滤。变量没设 → 策略为 false → 一行都查不到。
|
|
4. 仓储层强制 ws_id 必填，双保险：即使 RLS 被误关，代码也不会跨租户查。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
|
|
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
from sqlalchemy.pool import NullPool
|
|
|
|
from .config import get_settings
|
|
from .errors import AppError, ErrorCode
|
|
|
|
_settings = get_settings()
|
|
|
|
_engine_kwargs: dict[str, Any] = {"echo": _settings.db_echo, "pool_pre_ping": True}
|
|
if _settings.env == "test":
|
|
# pytest-asyncio 每个用例一个事件循环，连接池跨循环复用会报
|
|
# "attached to a different loop"，测试环境干脆不池化。
|
|
_engine_kwargs["poolclass"] = NullPool
|
|
else:
|
|
_engine_kwargs["pool_size"] = _settings.db_pool_size
|
|
|
|
engine = create_async_engine(_settings.database_url, **_engine_kwargs)
|
|
SessionFactory = async_sessionmaker(engine, expire_on_commit=False, autoflush=False)
|
|
|
|
|
|
async def get_session() -> AsyncIterator[AsyncSession]:
|
|
"""一个请求 = 一个事务。异常即整体回滚。"""
|
|
async with SessionFactory() as session:
|
|
async with session.begin():
|
|
yield session
|
|
|
|
|
|
async def bind_tenant(session: AsyncSession, ws_id: uuid.UUID, user_id: uuid.UUID | None) -> None:
|
|
"""把租户上下文写进当前事务。必须在任何业务查询之前调用。"""
|
|
await session.execute(
|
|
text("SELECT set_config('app.ws_id', :ws, true)"), {"ws": str(ws_id)}
|
|
)
|
|
if user_id is not None:
|
|
await session.execute(
|
|
text("SELECT set_config('app.user_id', :uid, true)"), {"uid": str(user_id)}
|
|
)
|
|
|
|
|
|
async def bind_user(session: AsyncSession, user_id: uuid.UUID) -> None:
|
|
"""只绑用户（登录后尚未选定工作区的场景）。"""
|
|
await session.execute(
|
|
text("SELECT set_config('app.user_id', :uid, true)"), {"uid": str(user_id)}
|
|
)
|
|
|
|
|
|
async def current_app_role(session: AsyncSession) -> dict[str, Any]:
|
|
"""自检用：确认应用连接的不是特权角色。"""
|
|
row = (
|
|
await session.execute(
|
|
text(
|
|
"SELECT current_user AS role, "
|
|
"pg_has_role(current_user, 'pg_read_all_data', 'member') AS read_all, "
|
|
"(SELECT rolsuper FROM pg_roles WHERE rolname = current_user) AS is_super, "
|
|
"(SELECT rolbypassrls FROM pg_roles WHERE rolname = current_user) AS bypass_rls"
|
|
)
|
|
)
|
|
).mappings().one()
|
|
return dict(row)
|
|
|
|
|
|
class Repo:
|
|
"""仓储基类：ws_id 必填，且没有默认值——不给"忘了传"留余地。"""
|
|
|
|
def __init__(self, session: AsyncSession, ws_id: uuid.UUID) -> None:
|
|
if ws_id is None:
|
|
raise AppError(ErrorCode.BAD_REQUEST, "缺少工作区上下文")
|
|
self.session = session
|
|
self.ws_id = ws_id
|
|
|
|
|
|
class UserRepo:
|
|
"""全局仓储（app_user 无租户），仅认证流程使用。"""
|
|
|
|
def __init__(self, session: AsyncSession) -> None:
|
|
self.session = session
|