"""会话内不可写迁移模板与按例独立数据库；不缓存跨次运行的数据库。 管理连接只来自显式隔离配置。所有清理同时核验随机名称、数据库 OID、 拥有者与本次会话注记，不按名称前缀扫描数据库。 """ from __future__ import annotations import fcntl import hashlib import json import os import time import uuid from collections.abc import Callable, Iterator from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path from typing import Any import psycopg from psycopg import sql from psycopg.conninfo import conninfo_to_dict, make_conninfo from muse.共享.调用身份 import 用途 from muse.基础设施.数据库.迁移 import 执行迁移 from muse.基础设施.数据库.连接 import 数据库工厂 from muse.配置 import 数据库引用 角色 = {用途.维护: "muse_maint", 用途.生产: "muse_app", 用途.评测: "muse_eval"} # 系统种子表登记：迁移 SQL 或导入种子写入、共享库例间清空必须保留的表。 # 新增迁移若 INSERT 种子数据，必须同步登记；生命周期单测会机械拦截漏登记。 种子表 = frozenset({ ("public", "muse_migration"), # 迁移账本，由迁移执行器维护 ("public", "muse_method_index_guard"), # V0059 检索索引代次单例 ("public", "muse_raw_namespace"), # V0004 原文命名空间单例 ("evaluation", "muse_raw_namespace"), # V0004 原文命名空间单例（evaluation 库） ("metadata", "type_registry"), # 内置结构种子，由导入种子写入 ("metadata", "schema_version"), # 内置结构种子，由导入种子写入 ("metadata", "work_extension"), # 内置结构种子，由导入种子写入 }) @dataclass(frozen=True) class 库身份: 名称: str oid: int owner: int 注记: str class PostgreSQL后端: """真实管理操作；离线行为测试只替换这一环境边界。""" def __init__(self, 管理URL: str, 共享目录: Path): self.管理URL = 管理URL self.共享目录 = 共享目录 def 环境身份(self) -> dict: with psycopg.connect(self.管理URL, autocommit=True) as 连: return { "major": 连.info.server_version // 10000, "extensions": 连.execute( "SELECT extname, extversion FROM pg_extension ORDER BY extname" ).fetchall(), } def 初始化角色(self, 定义: str) -> bool: # 同一 pytest run 的各 worker 共用临时标记，数据库锁保护角色的创建。 # 标记只缓存本次执行，键不包含明文连接串。 self.共享目录.mkdir(mode=0o700, parents=True, exist_ok=True) 键 = hashlib.sha256((self.管理URL + 定义).encode()).hexdigest() with (self.共享目录 / 键).open("a+") as 标记: fcntl.flock(标记, fcntl.LOCK_EX) 标记.seek(0) if 标记.read() == "ready": return False with psycopg.connect(self.管理URL, autocommit=True) as 连: 连.execute("SELECT pg_advisory_lock(7050191705)") try: 连.execute(定义) finally: 连.execute("SELECT pg_advisory_unlock(7050191705)") 标记.write("ready") 标记.flush() return True def _读取身份(self, 连, 名称: str) -> tuple | None: return 连.execute( "SELECT oid, datdba, shobj_description(oid, 'pg_database'), datallowconn " "FROM pg_database WHERE datname = %s", (名称,), ).fetchone() def 创建(self, 名称: str, 注记: str, 模板: 库身份 | None = None) -> 库身份: with psycopg.connect(self.管理URL, autocommit=True) as 连: if 模板: self._核验(连, 模板, 模板=True) 连.execute( sql.SQL("CREATE DATABASE {} OWNER muse_maint TEMPLATE {}").format( sql.Identifier(名称), sql.Identifier(模板.名称 if 模板 else "template0") ) ) # 记录 OID 后再写注记；写失败仅清理刚刚确认创建的同一 OID。 行 = self._读取身份(连, 名称) assert 行 is not None try: 连.execute( sql.SQL("COMMENT ON DATABASE {} IS {}").format( sql.Identifier(名称), sql.Literal(注记) ) ) except BaseException: 当前 = self._读取身份(连, 名称) if 当前 and 当前[:2] == 行[:2]: 连.execute( sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(名称)) ) raise return 库身份(名称, 行[0], 行[1], 注记) def _客户端占用(self, 连, oid: int) -> int: """只数客户端会话：自动清理等后台进程不占连接，也不阻止建立克隆。""" return 连.execute( "SELECT count(*) FROM pg_stat_activity " "WHERE datid = %s AND backend_type = 'client backend'", (oid,), ).fetchone()[0] def _核验(self, 连, 身份: 库身份, *, 模板: bool = False) -> None: 行 = self._读取身份(连, 身份.名称) if 行 is None or 行[:3] != (身份.oid, 身份.owner, 身份.注记): raise RuntimeError(f"数据库身份已改变，拒绝操作：{身份.名称}") if not 模板: return # 刚关闭的会话行可能仍在回收；有界等待真实客户端会话退出后再判定。 截止 = time.monotonic() + 10.0 while 行[3] or self._客户端占用(连, 身份.oid): if time.monotonic() >= 截止: raise RuntimeError(f"模板未封存或仍有连接：{身份.名称}") time.sleep(0.05) 行 = self._读取身份(连, 身份.名称) if 行 is None or 行[:3] != (身份.oid, 身份.owner, 身份.注记): raise RuntimeError(f"数据库身份已改变，拒绝操作：{身份.名称}") def 封存(self, 身份: 库身份) -> None: with psycopg.connect(self.管理URL, autocommit=True) as 连: self._核验(连, 身份) 连.execute( sql.SQL("ALTER DATABASE {} ALLOW_CONNECTIONS false").format( sql.Identifier(身份.名称) ) ) self._核验(连, 身份, 模板=True) def 删除(self, 身份: 库身份) -> None: with psycopg.connect(self.管理URL, autocommit=True) as 连: self._核验(连, 身份) 连.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(身份.名称))) def 工厂(self, 库名: str | None, 目录: Path) -> dict[用途, 数据库工厂]: 目录.mkdir(parents=True, exist_ok=True) 结果 = {} for 用途值, 角色名 in 角色.items(): 参数 = conninfo_to_dict(self.管理URL) 参数["user"] = 角色名 if 库名 is not None: 参数["dbname"] = 库名 路径 = 目录 / f"{用途值.value}.txt" # 原子创建为 0600，避免明文连接引用曾以默认权限短暂出现。 with open( 路径, "x", encoding="utf-8", opener=lambda path, flags: os.open(path, flags, 0o600) ) as 文件: 文件.write(make_conninfo(**参数)) 结果[用途值] = 数据库工厂(数据库引用("受控存储", str(路径)), 用途值) return 结果 def 迁移(self, 身份: 库身份, 目录: Path, 引用目录: Path) -> int: 工厂 = self.工厂(身份.名称, 引用目录) with 工厂[用途.维护].连接() as 连: return len(执行迁移(连, 目录)) def 清空业务数据(self, 身份: 库身份) -> None: """例间清空业务表并重置序列；保留 种子表 登记的系统种子，比整库销毁快百倍。""" 参数 = conninfo_to_dict(self.管理URL) 参数["user"] = 角色[用途.维护] 参数["dbname"] = 身份.名称 with psycopg.connect(make_conninfo(**参数), autocommit=True) as 连: # 终止同一测试残留的空闲会话，避免 TRUNCATE 产生锁等待 连.execute( """ SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = %s AND pid != pg_backend_pid() AND backend_type = 'client backend' AND state IN ('idle', 'idle in transaction', 'idle in transaction (aborted)'); """, (身份.名称,), ) # 拉出全部非系统 schema 的基表，在 Python 侧按 种子表 过滤； # 未知 schema 与未来新增的运行时表默认清空，只有登记过的种子保留。 表行 = 连.execute( """ SELECT table_schema, table_name FROM information_schema.tables WHERE table_type = 'BASE TABLE' AND table_schema NOT IN ('pg_catalog', 'information_schema') ORDER BY table_schema, table_name; """ ).fetchall() 目标 = [f'"{s}"."{t}"' for s, t in 表行 if (s, t) not in 种子表] if 目标: 连.execute("TRUNCATE TABLE " + ", ".join(目标) + " RESTART IDENTITY CASCADE") class 数据库底座: """每个 session/worker 一个实例；模板构建失败不得进入可复用集合。""" def __init__(self, 后端: Any, 根: Path, 临时目录: Path, *, worker: str = "main"): self.后端 = 后端 self.根 = 根 self.临时目录 = 临时目录 self.会话 = uuid.uuid4().hex self.worker = worker self.模板: dict[str, 库身份] = {} self.共享: dict[str, 库身份] = {} self.存活: dict[str, 库身份] = {} self.事件: list[dict] = [] self._已关闭 = False self._角色已初始化 = False self._环境: dict | None = None self._输入: dict[str, bytes] = {} self._角色定义 = (根 / "数据库/初始化/用途角色.sql").read_bytes() self._输入 = {p.name: p.read_bytes() for p in sorted((根 / "数据库/迁移").glob("*.sql"))} if not self._输入: raise ValueError("迁移目录为空") def 记录(self, 动作: str, **字段) -> None: self.事件.append({"action": 动作, "time": time.time(), "worker": self.worker, **字段}) def 初始化(self) -> None: if self._已关闭: raise RuntimeError("数据库底座已关闭，拒绝新借用") if not self._角色已初始化: 已执行 = self.后端.初始化角色(self._角色定义.decode()) self._环境 = self.后端.环境身份() self._角色已初始化 = True self.记录("role_bootstrap" if 已执行 else "role_bootstrap_reuse") def 端点(self, 目录: Path): self.初始化() return self.后端.工厂(None, 目录) def _新库(self, 模板: 库身份 | None = None) -> 库身份: self.初始化() 名称 = f"muse_test_{self.会话[:12]}_{uuid.uuid4().hex[:12]}" try: 身份 = self.后端.创建(名称, self.会话, 模板) except BaseException as 错误: # 创建已到达服务器而回包丢失时，不猜测所有权后删除；保存精确名称供核对。 self.记录("create_failed", database=名称, reason=type(错误).__name__) raise self.存活[名称] = 身份 self.记录("clone" if 模板 else "empty_create", database=名称) return 身份 def _删除(self, 身份: 库身份) -> None: if self.存活.get(身份.名称) != 身份: raise RuntimeError("拒绝删除本会话未登记的数据库") try: self.后端.删除(身份) except BaseException as 错误: self.记录("cleanup_failed", database=身份.名称, reason=type(错误).__name__) raise else: self.存活.pop(身份.名称) self.记录("cleanup", database=身份.名称) def _指纹( self, 版本: tuple[int, ...] | None, 种子指纹: str = "" ) -> tuple[str, dict[str, bytes]]: self.初始化() # 输入冻结在会话开始；源码在运行中变化时拒绝，而非悄悄复用旧模板。 当前 = {p.name: p.read_bytes() for p in sorted((self.根 / "数据库/迁移").glob("*.sql"))} if ( 当前 != self._输入 or (self.根 / "数据库/初始化/用途角色.sql").read_bytes() != self._角色定义 ): raise RuntimeError("会话中迁移或角色定义已改变，请重启测试会话") 输入 = { 名: 内容 for 名, 内容 in self._输入.items() if 版本 is None or int(名.split("__")[0][1:]) in 版本 } if 版本 is not None and {int(名.split("__")[0][1:]) for 名 in 输入} != set(版本): raise ValueError("请求的迁移版本不存在") 指纹 = hashlib.sha256( json.dumps( { "environment": self._环境, "roles": hashlib.sha256(self._角色定义).hexdigest(), "migrations": [ (名, hashlib.sha256(内容).hexdigest()) for 名, 内容 in 输入.items() ], "seed": 种子指纹, }, sort_keys=True, ).encode() ).hexdigest() return 指纹, 输入 def 基线( self, 版本: tuple[int, ...] | None = None, *, 种子指纹: str = "", 导入种子: Callable | None = None, ) -> 库身份: if bool(种子指纹) != bool(导入种子): raise ValueError("种子导入函数与内容指纹必须同时提供") 指纹, 输入 = self._指纹(版本, 种子指纹) if 指纹 in self.模板: return self.模板[指纹] 父库 = self.基线(版本) if 导入种子 else None 身份 = self._新库(父库) try: 目录 = self.临时目录 / (指纹 + "_" + uuid.uuid4().hex[:8]) 目录.mkdir(parents=True) if 导入种子: 导入种子(self.后端.工厂(身份.名称, 目录 / "引用")) self.记录("seed_import", database=身份.名称) else: SQL目录 = 目录 / "迁移" SQL目录.mkdir() for 名, 内容 in 输入.items(): (SQL目录 / 名).write_bytes(内容) 数量 = self.后端.迁移(身份, SQL目录, 目录 / "引用") self.记录("ddl", database=身份.名称, files=数量) self.后端.封存(身份) self.模板[指纹] = 身份 self.记录("template_ready", database=身份.名称, fingerprint=指纹) return 身份 except BaseException: self._删除(身份) raise @contextmanager def 借库( self, 目录: Path, 版本: tuple[int, ...] | None = None, *, 空库: bool = False, 种子指纹: str = "", 导入种子: Callable | None = None, ) -> Iterator[dict]: 身份 = self._新库(None if 空库 else self.基线(版本, 种子指纹=种子指纹, 导入种子=导入种子)) try: yield self.后端.工厂(身份.名称, 目录) finally: self._删除(身份) @contextmanager def 共享库( self, 目录: Path, 版本: tuple[int, ...] | None = None, *, 种子指纹: str = "", 导入种子: Callable | None = None, 清空: bool = True, ) -> Iterator[dict]: """会话内复用同一已迁移库，例后清空业务数据；仅会话结束时统一销毁。""" self.初始化() 指纹, _ = self._指纹(版本, 种子指纹) if 指纹 not in self.共享: 身份 = self._新库(self.基线(版本, 种子指纹=种子指纹, 导入种子=导入种子)) self.共享[指纹] = 身份 self.记录("shared_create", database=身份.名称, fingerprint=指纹) else: 身份 = self.共享[指纹] # 兜底：上例若被强杀（kill -9）例后清空未执行；借用前幂等清一次。 # 首次借用是刚迁移完的干净库，无需清。 if 清空 and hasattr(self.后端, "清空业务数据"): try: self.后端.清空业务数据(身份) self.记录("shared_clean_before", database=身份.名称) except Exception as 错误: self.记录( "shared_clean_before_failed", database=身份.名称, reason=type(错误).__name__, ) raise self.记录("shared_reuse", database=身份.名称, fingerprint=指纹) 工厂 = self.后端.工厂(身份.名称, 目录) try: yield 工厂 finally: if 清空 and hasattr(self.后端, "清空业务数据"): try: self.后端.清空业务数据(身份) self.记录("shared_clean", database=身份.名称) except Exception as 错误: self.记录( "shared_clean_failed", database=身份.名称, reason=type(错误).__name__ ) raise def 关闭(self) -> None: self._已关闭 = True 错误 = [] for 身份 in reversed(tuple(self.存活.values())): try: self._删除(身份) except Exception as exc: 错误.append(exc) if 错误: raise ExceptionGroup("本会话测试数据库清理失败，见成本事件的具体数据库名称", 错误)