Steven
2b493bd9f8
CI / backend (push) Canceled after 0s
跑通迁移与测试后修掉的真问题： - alembic 0001：asyncpg 走预备语句，一次 execute 只能一条 SQL。 CREATE FUNCTION / CREATE POLICY / GRANT 全部拆成单语句逐条执行， 否则报 "cannot insert multiple commands into a prepared statement"。 - 创建工作区：ORM flush 只按 relationship 排序，裸外键不算数， member 可能先于 workspace 落地撞 FK。改为先 flush workspace 再插 member。 - datetime.UTC 是 3.11+ 才有的，本机 3.10 直接 ImportError。改用 timezone.utc。 - 幂等中间件：原来会缓存 4xx。校验失败被缓存后，用户改对参数仍拿到同一个 422。 改为只缓存 <400。 - 发码节流窗口改为可配置 SMS_THROTTLE_SECONDS（测试置 0）， 节流用例自己显式打开，避免"整库关掉节流"导致节流从未被验证。 - 租户隔离用例：pytest.raises(Exception) 改为精确的 ProgrammingError， 否则任何异常都会让"应用角色不能建表"这条断言假通过。 验收：pytest 33 passed；ruff check / format 通过；真机跑 uvicorn， /health、/version、登录、建工作区、列表、/auth/me 全通。
96 lines
3.6 KiB
Python
96 lines
3.6 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, 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
|