403 lines
17 KiB
Python
403 lines
17 KiB
Python
"""会话内不可写迁移模板与按例独立数据库；不缓存跨次运行的数据库。
|
||
|
||
管理连接只来自显式隔离配置。所有清理同时核验随机名称、数据库 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"}
|
||
|
||
|
||
@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)');
|
||
""",
|
||
(身份.名称,),
|
||
)
|
||
表行 = 连.execute(
|
||
"""
|
||
SELECT table_schema, table_name
|
||
FROM information_schema.tables
|
||
WHERE table_type = 'BASE TABLE'
|
||
AND (
|
||
(table_schema IN ('public', 'evaluation', 'oracle')
|
||
AND table_name NOT IN ('muse_migration', 'muse_method_index_guard', 'muse_raw_namespace'))
|
||
OR
|
||
(table_schema = 'metadata'
|
||
AND table_name IN ('schema_binding', 'policy_version', 'activation_receipt'))
|
||
)
|
||
ORDER BY table_schema, table_name;
|
||
"""
|
||
).fetchall()
|
||
if 表行:
|
||
目标 = ", ".join(f'"{s}"."{t}"' for s, t in 表行)
|
||
连.execute(f"TRUNCATE TABLE {目标} 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.共享[指纹]
|
||
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("本会话测试数据库清理失败，见成本事件的具体数据库名称", 错误)
|