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())
|