135 lines
4.0 KiB
Python
135 lines
4.0 KiB
Python
"""
|
|
PostgreSQL 运行账连接与迁移。
|
|
|
|
本包职责：解析 PG_DSN、安全拒绝现网库名、执行版本化 SQL。
|
|
禁止：智能体经本包写 MySQL 工单表；禁止 FLUSH / 删现网库。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import re
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
from urllib.parse import urlparse
|
|
|
|
from agent.config import Settings, get_settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# 现网运行账库名：新系统迁移脚本绝对不能碰
|
|
_FORBIDDEN_DB_NAMES = frozenset({"ytd_runtime"})
|
|
|
|
|
|
def migrations_dir() -> Path:
|
|
"""返回 inquiry-agent/migrations 目录。"""
|
|
return Path(__file__).resolve().parents[2] / "migrations"
|
|
|
|
|
|
def assert_safe_database(dsn: str) -> str:
|
|
"""
|
|
从 DSN 解析库名并校验。
|
|
|
|
返回库名；若为 ytd_runtime 则抛错。
|
|
"""
|
|
if not dsn:
|
|
raise RuntimeError("PG_DSN 为空，无法迁移")
|
|
parsed = urlparse(dsn)
|
|
db = (parsed.path or "").lstrip("/").split("?")[0]
|
|
if not db:
|
|
raise RuntimeError(f"无法从 DSN 解析库名: {dsn[:32]}...")
|
|
if db in _FORBIDDEN_DB_NAMES:
|
|
raise RuntimeError(f"禁止对现网库执行迁移: {db}")
|
|
return db
|
|
|
|
|
|
def connect_psycopg(dsn: str):
|
|
"""
|
|
打开 psycopg 连接（调用方负责 close）。
|
|
|
|
副作用：网络连库。
|
|
"""
|
|
import psycopg
|
|
|
|
return psycopg.connect(dsn, connect_timeout=5)
|
|
|
|
|
|
def apply_migrations(settings: Optional[Settings] = None) -> list[str]:
|
|
"""
|
|
按文件名顺序执行尚未应用的 *.sql。
|
|
|
|
返回：本次新执行的 version 列表。
|
|
无 DSN：抛错（自检可用 --dry-list）。
|
|
"""
|
|
cfg = settings or get_settings()
|
|
dsn = cfg.resolve_pg_dsn()
|
|
db_name = assert_safe_database(dsn)
|
|
logger.info("准备迁移 database=%s", db_name)
|
|
|
|
sql_files = sorted(migrations_dir().glob("*.sql"))
|
|
if not sql_files:
|
|
logger.warning("migrations 目录无 .sql 文件")
|
|
return []
|
|
|
|
applied: list[str] = []
|
|
with connect_psycopg(dsn) as conn:
|
|
conn.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS schema_migrations (
|
|
version TEXT PRIMARY KEY,
|
|
applied_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
|
description TEXT NOT NULL DEFAULT ''
|
|
)
|
|
"""
|
|
)
|
|
conn.commit()
|
|
for path in sql_files:
|
|
version = path.stem
|
|
row = conn.execute(
|
|
"SELECT 1 FROM schema_migrations WHERE version = %s",
|
|
(version,),
|
|
).fetchone()
|
|
if row:
|
|
logger.info("跳过已应用 migration=%s", version)
|
|
continue
|
|
sql = path.read_text(encoding="utf-8")
|
|
# 去掉纯注释行后的空脚本保护
|
|
body = "\n".join(
|
|
ln for ln in sql.splitlines() if not ln.strip().startswith("--")
|
|
).strip()
|
|
if not body:
|
|
continue
|
|
logger.info("执行 migration=%s", version)
|
|
conn.execute(sql)
|
|
# 若脚本未自行写入 schema_migrations，补写
|
|
exists = conn.execute(
|
|
"SELECT 1 FROM schema_migrations WHERE version = %s",
|
|
(version,),
|
|
).fetchone()
|
|
if not exists:
|
|
conn.execute(
|
|
"INSERT INTO schema_migrations(version, description) VALUES (%s, %s)",
|
|
(version, path.name),
|
|
)
|
|
conn.commit()
|
|
applied.append(version)
|
|
return applied
|
|
|
|
|
|
def list_migration_files() -> list[str]:
|
|
"""列出迁移文件名（不连库）。"""
|
|
return [p.name for p in sorted(migrations_dir().glob("*.sql"))]
|
|
|
|
|
|
# 会话 / 停等仓储壳（表已在 001；本包只提供读写空实现）
|
|
from agent.runtime_db.session_store import ( # noqa: E402
|
|
MemorySessionStore,
|
|
SessionRecord,
|
|
create_session_store,
|
|
)
|
|
from agent.runtime_db.suspend_store import ( # noqa: E402
|
|
MemorySuspendStore,
|
|
SuspendRecord,
|
|
create_suspend_store,
|
|
)
|