Steven
d390b7b58a
CI / backend (push) Canceled after 0s
fnsvr-01 的 PostgreSQL 在 8432 端口（不是 5432），实测 PostgreSQL 17.6 / Alpine。 已在该实例建好独立库与角色： - dadaotask（开发）/ dadaotask_test（测试），owner 都是 dadao_owner - dadao_owner 迁移用；dadao_app 应用用（NOSUPERUSER NOBYPASSRLS） - 两个库都已 ALTER SCHEMA public OWNER TO dadao_owner + REVOKE CREATE FROM PUBLIC + GRANT 给 dadao_app 代码改动： - config 增加 test_database_url / test_migration_database_url，并给出 resolved_database_url() / resolved_migration_database_url()：ENV=test 且有测试库时优先用测试库， 否则退回开发库（CI 里直接给 DATABASE_URL 即可，行为不变）。 - db.py / alembic/env.py / tests/conftest.py 统一走这两个解析方法。 原先 conftest 清表用的是 settings.migration_database_url，而应用连的是测试库， 一旦测试库与开发库不同，TRUNCATE 就清到另一个库上——表现为"测试库脏数据越攒越多、 列表类断言随机失败"。这是本次真踩到的坑，不是假想。 - tests/conftest.py 保留 E402 豁免（环境变量必须在 import app 模块前写好）。 验证： - 远程 dadaotask_test 上 pytest 33 passed，跑完表内 0 行（确认清表打在同一个库） - 远程 dadaotask 上 alembic upgrade head 成功，/health 与 /version 正常，db_role 为 dadao_app 且三项全 false - 本机便携 PG 16.4 上同样 33 passed（两个 PG 大版本都过） - ruff check / format 干净 已知限制：该实例没有 pgvector（只有 pg_trgm / pgcrypto），M3 向量检索前必须先定扩展来源， 已写进 2.Design/DEVELOPMENT.md。
99 lines
3.7 KiB
Python
99 lines
3.7 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()
|
|
|
|
# 测试环境优先用独立测试库：测试会 TRUNCATE 所有表，落到开发库上数据就没了。
|
|
_db_url = _settings.resolved_database_url()
|
|
|
|
_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(_db_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
|