382 lines
13 KiB
Python
382 lines
13 KiB
Python
"""CLI 命令分发：argparse 子命令 → 调用 service 层。
|
||
|
||
v1.1.1：11 个子命令（在 9 个基础上新增）
|
||
- collect-from-platform：Playwright RPA 抓 SmartCampus HR 平台数据
|
||
- parse-materials：解析已下载的材料 PDF/Word → 文本
|
||
|
||
原有 9 个：
|
||
- health / version / config-path / init-db
|
||
- scan-expiring-contracts / generate-renewal-task
|
||
- dispatch-task-notice / mark-contract-renewed
|
||
- render-html
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import sys
|
||
from typing import Callable, List, Optional
|
||
|
||
from db.connection import get_connection, init_db
|
||
from db import repository as repo
|
||
from service import dispatcher, html_renderer, material_parser, platform_collector, scanner, task_generator
|
||
from util.constants import DEFAULT_THRESHOLDS, HUMAN_WAIT_TIMEOUT, SKILL_VERSION
|
||
from util.logging_setup import get_skill_logger, setup_skill_logging
|
||
from util.runtime_paths import (
|
||
get_db_path,
|
||
get_dispatch_dir,
|
||
get_env_path,
|
||
get_output_dir,
|
||
)
|
||
|
||
logger = get_skill_logger()
|
||
|
||
|
||
# ---------- 命令实现 ----------
|
||
|
||
def cmd_health(_args: argparse.Namespace) -> int:
|
||
"""健康检查。"""
|
||
setup_skill_logging()
|
||
db_path = get_db_path()
|
||
db_exists = db_path.exists()
|
||
db_writable = False
|
||
error_msg = None
|
||
|
||
try:
|
||
conn = get_connection()
|
||
try:
|
||
conn.execute("SELECT 1").fetchone()
|
||
db_writable = True
|
||
finally:
|
||
conn.close()
|
||
except Exception as e:
|
||
error_msg = str(e)
|
||
|
||
output_dir = get_output_dir()
|
||
dispatch_dir = get_dispatch_dir()
|
||
checks = [
|
||
("python_version", f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}", True),
|
||
("db_path", str(db_path), True),
|
||
("db_exists", str(db_exists), db_exists),
|
||
("db_writable", str(db_writable), db_writable),
|
||
("output_dir", str(output_dir), True),
|
||
("dispatch_dir", str(dispatch_dir), True),
|
||
]
|
||
|
||
if db_exists and db_writable:
|
||
try:
|
||
conn = get_connection()
|
||
try:
|
||
counts = {
|
||
"contracts": conn.execute("SELECT COUNT(*) AS n FROM contracts").fetchone()["n"],
|
||
"performance_records": conn.execute("SELECT COUNT(*) AS n FROM performance_records").fetchone()["n"],
|
||
"alerts": conn.execute("SELECT COUNT(*) AS n FROM alerts").fetchone()["n"],
|
||
"renewal_tasks": conn.execute("SELECT COUNT(*) AS n FROM renewal_tasks").fetchone()["n"],
|
||
"dispatch_log": conn.execute("SELECT COUNT(*) AS n FROM dispatch_log").fetchone()["n"],
|
||
}
|
||
finally:
|
||
conn.close()
|
||
checks.append(("seed_loaded", f"contracts={counts['contracts']} perf={counts['performance_records']}", counts["contracts"] > 0))
|
||
except Exception as e:
|
||
checks.append(("seed_check", f"ERROR: {e}", False))
|
||
|
||
all_ok = all(c[2] for c in checks)
|
||
for name, val, ok in checks:
|
||
marker = "OK " if ok else "FAIL"
|
||
print(f"[{marker}] {name}: {val}")
|
||
if error_msg:
|
||
print(f"[FAIL] error: {error_msg}")
|
||
|
||
if all_ok:
|
||
print("OK: health check passed")
|
||
return 0
|
||
print("ERROR: health check failed")
|
||
return 1
|
||
|
||
|
||
def cmd_version(_args: argparse.Namespace) -> int:
|
||
print(f"{SKILL_VERSION}")
|
||
return 0
|
||
|
||
|
||
def cmd_config_path(_args: argparse.Namespace) -> int:
|
||
print(get_env_path())
|
||
return 0
|
||
|
||
|
||
def cmd_init_db(_args: argparse.Namespace) -> int:
|
||
counts = init_db()
|
||
print(f"OK: init-db complete")
|
||
for k, v in counts.items():
|
||
print(f" {k}: {v} rows")
|
||
return 0
|
||
|
||
|
||
def cmd_scan(args: argparse.Namespace) -> int:
|
||
thresholds = _parse_thresholds(args.thresholds)
|
||
conn = get_connection()
|
||
try:
|
||
result = scanner.scan(conn, thresholds=thresholds)
|
||
finally:
|
||
conn.close()
|
||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||
return 0
|
||
|
||
|
||
def cmd_generate_task(args: argparse.Namespace) -> int:
|
||
contract_id = args.contract_id
|
||
task_id = args.task_id
|
||
conn = get_connection()
|
||
try:
|
||
result = task_generator.regenerate_task_for_contract_id(conn, contract_id, task_id)
|
||
conn.commit()
|
||
except ValueError as e:
|
||
print(f"ERROR: {e}")
|
||
return 2
|
||
finally:
|
||
conn.close()
|
||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||
return 0
|
||
|
||
|
||
def cmd_dispatch(args: argparse.Namespace) -> int:
|
||
conn = get_connection()
|
||
try:
|
||
result = dispatcher.dispatch(conn, args.unit_id, args.month or "")
|
||
finally:
|
||
conn.close()
|
||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||
return 0
|
||
|
||
|
||
def cmd_mark_renewed(args: argparse.Namespace) -> int:
|
||
conn = get_connection()
|
||
try:
|
||
from db import repository as repo
|
||
|
||
contract = repo.get_contract(conn, args.contract_id)
|
||
if contract is None:
|
||
print(f"ERROR: contract_id={args.contract_id!r} 不存在")
|
||
return 2
|
||
|
||
with conn:
|
||
n = repo.update_contract_status(conn, args.contract_id, "已续签", args.material_id)
|
||
closed = repo.close_open_alerts_for_contract(conn, args.contract_id)
|
||
print(json.dumps({
|
||
"contract_id": args.contract_id,
|
||
"status": "已续签",
|
||
"material_id": args.material_id,
|
||
"updated_rows": n,
|
||
"closed_alerts": closed,
|
||
}, ensure_ascii=False, indent=2))
|
||
finally:
|
||
conn.close()
|
||
return 0
|
||
|
||
|
||
def cmd_render_html(args: argparse.Namespace) -> int:
|
||
conn = get_connection()
|
||
try:
|
||
output_path = html_renderer.render_task_html(conn, args.task_id)
|
||
except ValueError as e:
|
||
print(f"ERROR: {e}")
|
||
return 2
|
||
except FileNotFoundError as e:
|
||
print(f"ERROR: {e}")
|
||
return 3
|
||
finally:
|
||
conn.close()
|
||
print(json.dumps({
|
||
"task_id": args.task_id,
|
||
"output_path": str(output_path),
|
||
"size_bytes": output_path.stat().st_size,
|
||
}, ensure_ascii=False, indent=2))
|
||
return 0
|
||
|
||
|
||
def cmd_collect_platform(args: argparse.Namespace) -> int:
|
||
"""Playwright RPA 抓 SmartCampus HR 平台数据 → 落库。
|
||
|
||
默认显示浏览器，以便登录态失效时可立即手工登录；
|
||
传 --headless 可复用已保存的 profile 进行无头采集。
|
||
|
||
退出码：
|
||
0 = 成功（即使抓到 0 行也返回 0，错误信息在 output["error"]）
|
||
1 = 通用异常
|
||
2 = RPA 结构化错误（LoginTimeout / MissingBrowser / TableNotFound 等）
|
||
"""
|
||
try:
|
||
result = platform_collector.collect_all(
|
||
limit=args.limit,
|
||
headless=args.headless,
|
||
human_wait_timeout=args.human_wait_timeout or HUMAN_WAIT_TIMEOUT,
|
||
)
|
||
except platform_collector.RPAError as e:
|
||
# v1.1.6：结构化错误（ERROR:LOGIN_TIMEOUT 等）→ exit 2
|
||
print(json.dumps({
|
||
"error": e.code,
|
||
"detail": str(e),
|
||
"fetched_rows": 0,
|
||
"saved_contracts": 0,
|
||
"downloaded_materials": 0,
|
||
"contracts": [],
|
||
}, ensure_ascii=False, indent=2))
|
||
return 2
|
||
except Exception as e:
|
||
print(json.dumps({"error": "exception", "detail": str(e)},
|
||
ensure_ascii=False, indent=2))
|
||
return 1
|
||
|
||
print(json.dumps(result, ensure_ascii=False, indent=2))
|
||
return 0
|
||
|
||
|
||
def cmd_parse_materials(args: argparse.Namespace) -> int:
|
||
"""对已落库的材料 PDF/Word 跑文本提取（pypdf / python-docx / OCR）。
|
||
|
||
--contract-id C001: 仅处理该合同
|
||
不传：处理所有 parse_method='pending' 或 extracted_text IS NULL 的材料
|
||
"""
|
||
conn = get_connection()
|
||
try:
|
||
if args.contract_id:
|
||
rows = repo.list_materials_for_contract(conn, args.contract_id)
|
||
else:
|
||
rows = repo.list_pending_materials(conn)
|
||
|
||
results: list[dict] = []
|
||
for r in rows:
|
||
from pathlib import Path
|
||
fp = Path(r["file_path"])
|
||
text, method = material_parser.extract_text(fp)
|
||
if text:
|
||
repo.insert_material(
|
||
conn,
|
||
r["contract_id"], r["file_name"], r["file_path"],
|
||
text, method,
|
||
)
|
||
else:
|
||
repo.insert_material(
|
||
conn,
|
||
r["contract_id"], r["file_name"], r["file_path"],
|
||
"", "failed",
|
||
)
|
||
results.append({
|
||
"contract_id": r["contract_id"],
|
||
"file_name": r["file_name"],
|
||
"parse_method": method,
|
||
"text_length": len(text),
|
||
})
|
||
|
||
conn.commit()
|
||
print(json.dumps({
|
||
"processed": len(results),
|
||
"results": results,
|
||
}, ensure_ascii=False, indent=2))
|
||
return 0
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
# ---------- argparse ----------
|
||
|
||
def _parse_thresholds(s: str) -> tuple[int, int, int]:
|
||
parts = [int(p.strip()) for p in s.split(",")]
|
||
if len(parts) != 3 or sorted(parts, reverse=True) != parts:
|
||
raise argparse.ArgumentTypeError(
|
||
f"thresholds 必须是 3 个整数（黄,橙,红）且从大到小，例如 90,60,30；收到 {s!r}"
|
||
)
|
||
return (parts[0], parts[1], parts[2])
|
||
|
||
|
||
def build_parser() -> argparse.ArgumentParser:
|
||
p = argparse.ArgumentParser(
|
||
prog="main.py",
|
||
description=f"合同到期智能预警与续聘辅助 v{SKILL_VERSION}",
|
||
)
|
||
sub = p.add_subparsers(dest="cmd", required=True)
|
||
|
||
sub.add_parser("health", help="健康检查（Python / DB / 目录）").set_defaults(handler=cmd_health)
|
||
sub.add_parser("version", help="查看版本").set_defaults(handler=cmd_version)
|
||
sub.add_parser("config-path", help="查看 .env 路径").set_defaults(handler=cmd_config_path)
|
||
sub.add_parser("init-db", help="建表 + 加载样例数据").set_defaults(handler=cmd_init_db)
|
||
|
||
sp = sub.add_parser(
|
||
"scan-expiring-contracts", help="扫描即将到期合同（async/cron）",
|
||
)
|
||
sp.add_argument(
|
||
"--thresholds", default="90,60,30",
|
||
help="三级预警阈值（黄,橙,红）从大到小，默认 90,60,30",
|
||
)
|
||
sp.set_defaults(handler=cmd_scan)
|
||
|
||
sp = sub.add_parser(
|
||
"generate-renewal-task", help="为指定合同生成续聘任务书",
|
||
)
|
||
sp.add_argument("--contract-id", required=True, help="合同 ID")
|
||
sp.add_argument("--task-id", default=None, help="可选：指定 task_id（默认自动生成）")
|
||
sp.set_defaults(handler=cmd_generate_task)
|
||
|
||
sp = sub.add_parser(
|
||
"dispatch-task-notice", help="按单位/月份派发续聘通知（dry-run）",
|
||
)
|
||
sp.add_argument("--unit-id", required=True, help="二级单位 ID")
|
||
sp.add_argument("--month", default="", help="月份 YYYY-MM（留空 = 全部）")
|
||
sp.set_defaults(handler=cmd_dispatch)
|
||
|
||
sp = sub.add_parser(
|
||
"mark-contract-renewed", help="标记合同已续签",
|
||
)
|
||
sp.add_argument("--contract-id", required=True, help="合同 ID")
|
||
sp.add_argument("--material-id", required=True, help="续签材料 ID")
|
||
sp.set_defaults(handler=cmd_mark_renewed)
|
||
|
||
sp = sub.add_parser(
|
||
"render-html", help="渲染任务书 HTML",
|
||
)
|
||
sp.add_argument("--task-id", required=True, help="任务 ID")
|
||
sp.set_defaults(handler=cmd_render_html)
|
||
|
||
sp = sub.add_parser(
|
||
"collect-from-platform",
|
||
help="Playwright RPA 抓 SmartCampus HR 平台（合同表 + 材料包）",
|
||
)
|
||
sp.add_argument("--limit", type=int, default=0,
|
||
help="限制抓取行数（0=全量）")
|
||
browser_mode = sp.add_mutually_exclusive_group()
|
||
browser_mode.add_argument("--headless", action="store_true",
|
||
help="无头运行（仅适合已保存有效登录态的自动化任务）")
|
||
browser_mode.add_argument("--show-browser", dest="headless", action="store_false",
|
||
help=argparse.SUPPRESS)
|
||
sp.add_argument("--human-wait-timeout", type=int, default=None,
|
||
help="手工登录 fallback 等待秒数（默认 180 = HUMAN_WAIT_TIMEOUT）")
|
||
sp.set_defaults(handler=cmd_collect_platform)
|
||
|
||
sp = sub.add_parser(
|
||
"parse-materials", help="解析已下载的 PDF/Word 材料 → 文本",
|
||
)
|
||
sp.add_argument("--contract-id", default="",
|
||
help="仅解析该合同（留空=处理所有 pending）")
|
||
sp.set_defaults(handler=cmd_parse_materials)
|
||
|
||
return p
|
||
|
||
|
||
def main(argv: Optional[List[str]] = None) -> int:
|
||
argv = argv if argv is not None else sys.argv[1:]
|
||
if not argv:
|
||
print(f"合同到期智能预警与续聘辅助 v{SKILL_VERSION}")
|
||
print("Commands: health, version, config-path, init-db, "
|
||
"scan-expiring-contracts, generate-renewal-task, "
|
||
"dispatch-task-notice, mark-contract-renewed, render-html, "
|
||
"collect-from-platform, parse-materials")
|
||
return 1
|
||
|
||
setup_skill_logging()
|
||
|
||
parser = build_parser()
|
||
args = parser.parse_args(argv)
|
||
handler: Callable[[argparse.Namespace], int] = args.handler
|
||
return handler(args)
|