77 lines
2.5 KiB
Python
77 lines
2.5 KiB
Python
|
"""数据处理运行表的显式检查与安装命令。"""
|
|||
|
|
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import argparse
|
|||
|
|
from urllib.parse import urlsplit
|
|||
|
|
|
|||
|
|
from app.modules.data_process.store import DataProcessStore
|
|||
|
|
|
|||
|
|
REQUIRED_TASK_COLUMNS = (
|
|||
|
|
"generation_run_id",
|
|||
|
|
"results_confirmed",
|
|||
|
|
"workflow_step",
|
|||
|
|
"preview_status",
|
|||
|
|
"preview_progress",
|
|||
|
|
"preview_run_id",
|
|||
|
|
"preview_failure_reason",
|
|||
|
|
"preview_total_files",
|
|||
|
|
"preview_completed_files",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _target_label(database_url: str) -> str:
|
|||
|
|
parsed = urlsplit(database_url)
|
|||
|
|
database = parsed.path.strip("/") or "(unknown)"
|
|||
|
|
return f"{parsed.hostname or '(unknown)'}:{parsed.port or 5432}/{database}"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _schema_ready(store: DataProcessStore) -> bool:
|
|||
|
|
with store.connect() as conn:
|
|||
|
|
row = conn.execute(
|
|||
|
|
"""
|
|||
|
|
SELECT COUNT(*) = %s AS ready
|
|||
|
|
FROM information_schema.columns
|
|||
|
|
WHERE table_schema=current_schema()
|
|||
|
|
AND table_name='data_process_tasks'
|
|||
|
|
AND column_name = ANY(%s)
|
|||
|
|
""",
|
|||
|
|
(len(REQUIRED_TASK_COLUMNS), list(REQUIRED_TASK_COLUMNS)),
|
|||
|
|
).fetchone()
|
|||
|
|
return bool(row and row["ready"])
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main() -> int:
|
|||
|
|
parser = argparse.ArgumentParser(
|
|||
|
|
description="检查或显式安装数据处理运行表（不会由应用启动自动执行）"
|
|||
|
|
)
|
|||
|
|
action = parser.add_mutually_exclusive_group(required=True)
|
|||
|
|
action.add_argument("--check", action="store_true", help="只读检查迁移是否已安装")
|
|||
|
|
action.add_argument("--apply", action="store_true", help="执行 002 数据处理迁移")
|
|||
|
|
parser.add_argument(
|
|||
|
|
"--yes",
|
|||
|
|
action="store_true",
|
|||
|
|
help="确认允许修改 DATABASE_URL 指向的数据库；与 --apply 同时使用",
|
|||
|
|
)
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
store = DataProcessStore()
|
|||
|
|
target = _target_label(store.database_url)
|
|||
|
|
if args.check:
|
|||
|
|
ready = _schema_ready(store)
|
|||
|
|
print(f"数据处理 schema：{'已安装' if ready else '未安装'}；目标：{target}")
|
|||
|
|
return 0 if ready else 1
|
|||
|
|
if not args.yes:
|
|||
|
|
parser.error("--apply 必须同时提供 --yes，确认修改目标数据库")
|
|||
|
|
|
|||
|
|
print(f"正在安装数据处理 schema；目标：{target}")
|
|||
|
|
store.ensure_schema()
|
|||
|
|
if not _schema_ready(store):
|
|||
|
|
raise RuntimeError("迁移执行后仍未检测到 generation_run_id")
|
|||
|
|
print("数据处理 schema 安装完成")
|
|||
|
|
return 0
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
raise SystemExit(main())
|