"""Idempotently seed the workbench's original demo reservation tasks.""" from __future__ import annotations import argparse import json import sys from datetime import datetime, timedelta from pathlib import Path from uuid import UUID, uuid4 from zoneinfo import ZoneInfo from psycopg.types.json import Jsonb ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) from app.domain.schemas.workbench_runtime import WorkbenchRuntimeWorkOrder from app.native.wndb.core.connection import project_connection SCHEMA_SQL = ROOT / "resources/sql/011_workbench_work_orders.sql" SHANGHAI = ZoneInfo("Asia/Shanghai") DATASET_ID = "workbench-demo-orders@1" DEMO_DATE = datetime(2026, 9, 17, tzinfo=SHANGHAI) SEEDED_WORK_ORDER_PREFIX = "mock-work-order-" LEGACY_WORK_ORDER_SOURCES = ( "合成数据预约工单", "外围样板预约工单", ) WORK_ORDER_TEMPLATES = ( { "id": "incident-valve-dispatch", "code": "DEMO-WO-20260917-001", "title": "演示工单：华山路阀门边界确认", "summary": "异常事件分析后生成阀门边界确认任务，派发前需调度员复核影响范围。", "detail": "该条目为工作台演示数据，用于展示人工确认界面，不代表真实提交或岗位通知。", "recommendation": "确认影响用户、上下游压力和应急隔离边界后，再通知执行岗位到场。", "source": "异常事件分析", "location": "华山路沿线阀门组 VG-11", "dispatcher": "调度中心值班长", "assignee": "调度执行岗", "priority": "urgent", "risk_level": "attention", "phase": "dispatch", "start_hour": 10, "start_minute": 30, "duration_minutes": 40, "dispatch_lead_minutes": 30, "reply_window_minutes": 20, "reply_requirements": [ "确认阀门当前开度", "记录影响分区和用户范围", "回传派发确认截图", ], }, { "id": "valve-operation-execution", "code": "DEMO-WO-20260917-002", "title": "演示工单：北辰分区阀门开度调整", "summary": "根据管网智能调度方案，现场调整北辰分区边界阀门开度并复核压力。", "detail": "该工单模拟系统派发后的执行阶段，班组正在按方案对阀门要素进行操作。", "recommendation": "执行过程中保持调度语音联络，调整后立即复核末端压力和流量变化。", "source": "调度方案确认", "location": "北辰分区边界阀门 BV-07", "dispatcher": "管网调度岗", "assignee": "北辰抢修二组", "priority": "urgent", "risk_level": "attention", "phase": "execution", "start_hour": 11, "start_minute": 0, "duration_minutes": 55, "dispatch_lead_minutes": 25, "reply_window_minutes": 20, "reply_requirements": [ "上传阀门操作前后照片", "记录最终开度", "复核北辰末端压力不低于 0.22MPa", ], }, { "id": "pump-boundary-reply", "code": "DEMO-WO-20260917-003", "title": "演示工单：泵站边界复核复令", "summary": "泵站出水边界复核已完成，等待班组复令和调度确认闭环。", "detail": "该工单模拟执行完成后的复令阶段，用于检查处置结果、照片和指标回传。", "recommendation": "核对泵站出水压力、频率曲线和分区边界模拟结果后完成复令。", "source": "模型复核建议", "location": "北辰二级泵站 P-02", "dispatcher": "模型复核岗", "assignee": "泵站运行班组", "priority": "normal", "risk_level": "normal", "phase": "reply", "start_hour": 13, "start_minute": 0, "duration_minutes": 45, "dispatch_lead_minutes": 20, "reply_window_minutes": 35, "reply_requirements": [ "上传泵站运行截图", "填写出水压力复核结果", "确认模型边界已同步更新", ], }, ) def _build_activity(template: dict, scheduled_at: datetime, updated_at: datetime) -> list[dict]: created_at = scheduled_at - timedelta(minutes=template["dispatch_lead_minutes"]) activity = [{"sequence": 1, "occurred_at": created_at, "actor": template["dispatcher"], "kind": "created", "title": "演示工单创建", "detail": template["summary"]}] if template["phase"] != "dispatch": activity.append({"sequence": 2, "occurred_at": updated_at, "actor": template["assignee"], "kind": "updated", "title": "演示状态更新", "detail": template["detail"]}) return activity def build_demo_work_orders() -> list[dict]: """Build the three reservation tasks previously provided by the frontend mock.""" documents = [] for template in WORK_ORDER_TEMPLATES: scheduled_at = DEMO_DATE.replace( hour=template["start_hour"], minute=template["start_minute"], ) status = { "dispatch": "pending", "execution": "running", "reply": "awaiting_reply", }[template["phase"]] updated_at = { "dispatch": scheduled_at - timedelta(minutes=5), "execution": scheduled_at + timedelta(minutes=20), "reply": scheduled_at + timedelta(minutes=template["duration_minutes"]), }[template["phase"]] projected = { "id": f"{SEEDED_WORK_ORDER_PREFIX}{template['id']}", "code": template["code"], "scheduled_at": scheduled_at, "title": template["title"], "summary": template["summary"], "status": status, "risk_level": template["risk_level"], "updated_at": updated_at, "detail": template["detail"], "recommendation": template["recommendation"], "duration_minutes": template["duration_minutes"], "source": template["source"], "dispatcher": template["dispatcher"], "priority": template["priority"], "location": template["location"], "assignee": template["assignee"], "feedback_requirements": template["reply_requirements"], "activity": _build_activity(template, scheduled_at, updated_at), } documents.append( WorkbenchRuntimeWorkOrder.model_validate(projected).model_dump(mode="json") ) return documents def replace_demo_work_orders(conn, project_id: UUID, documents: list[dict]) -> None: """Remove the obsolete synthetic projection and replace this seed dataset.""" conn.execute( """ DELETE FROM workbench.work_orders WHERE project_id = %s AND ( work_order_id LIKE 'demo-work-order-%%' OR work_order_id LIKE %s OR document->>'source' = ANY(%s) ) """, ( project_id, f"{SEEDED_WORK_ORDER_PREFIX}%", list(LEGACY_WORK_ORDER_SOURCES), ), ) for document in documents: conn.execute( """ INSERT INTO workbench.work_orders (project_id, work_order_id, scheduled_at, updated_at, document) VALUES (%s, %s, %s, %s, %s) ON CONFLICT (project_id, work_order_id) DO UPDATE SET scheduled_at = EXCLUDED.scheduled_at, updated_at = EXCLUDED.updated_at, document = EXCLUDED.document """, ( project_id, document["id"], document["scheduled_at"], document["updated_at"], Jsonb(document), ), ) def publish_demo_work_order_batch_event(conn, project_id: UUID, count: int) -> None: """Invalidate the demo projection once, never once per seeded row.""" cursor = conn.execute( """INSERT INTO workbench.domain_event_log (project_id,event_id,topic,aggregate_type,aggregate_id, aggregate_revision,occurred_at,metadata) SELECT %s,%s,'workbench.work_order.changed','demo_work_order_dataset',%s, COALESCE(max(aggregate_revision),0)+1,now(),%s FROM workbench.domain_event_log WHERE project_id=%s AND aggregate_type='demo_work_order_dataset' AND aggregate_id=%s RETURNING event_seq""", ( project_id, uuid4(), DATASET_ID, Jsonb({"namespace": "demo-work-order", "count": count}), project_id, DATASET_ID, ), ) row = cursor.fetchone() conn.execute( "SELECT pg_notify('workbench_events', %s)", (f'{project_id}:{row["event_seq"]}',), ) def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("--database-name", default="tjwater_v2") parser.add_argument("--project-id", required=True, type=UUID) args = parser.parse_args() documents = build_demo_work_orders() with project_connection(args.database_name) as conn, conn.transaction(): conn.execute(SCHEMA_SQL.read_text(encoding="utf-8")) replace_demo_work_orders(conn, args.project_id, documents) publish_demo_work_order_batch_event(conn, args.project_id, len(documents)) print( json.dumps( { "database": args.database_name, "project_id": str(args.project_id), "dataset_id": DATASET_ID, "work_orders": [ { "id": document["id"], "code": document["code"], "title": document["title"], "status": document["status"], "feedback_requirement_count": len( document["feedback_requirements"] ), } for document in documents ], }, ensure_ascii=False, indent=2, ) ) return 0 if __name__ == "__main__": raise SystemExit(main())