270 lines
11 KiB
Python
270 lines
11 KiB
Python
"""Durable batch, execution and versioned workflow metadata; no device side effects."""
|
|
|
|
from contextlib import contextmanager
|
|
import json
|
|
import hashlib
|
|
from pathlib import Path
|
|
import re
|
|
import sqlite3
|
|
import time
|
|
import uuid
|
|
|
|
from artemis.config import DB_PATH
|
|
|
|
|
|
RECOVERY_ERRORS = {
|
|
"服务中断；未完成任务需手动重跑",
|
|
"服务中断，需手动重跑",
|
|
}
|
|
|
|
|
|
def identifier(value: str) -> str:
|
|
if not re.fullmatch(r"[A-Za-z0-9_-]{1,100}", value):
|
|
raise ValueError("无效的记录编号")
|
|
return value
|
|
|
|
|
|
def new_id() -> str:
|
|
return str(uuid.uuid4())
|
|
|
|
|
|
class AutomationStore:
|
|
def __init__(self, path: Path = DB_PATH):
|
|
self.path = Path(path)
|
|
|
|
@contextmanager
|
|
def connection(self):
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
conn = sqlite3.connect(self.path, timeout=15)
|
|
try:
|
|
conn.row_factory = sqlite3.Row
|
|
conn.executescript("""
|
|
CREATE TABLE IF NOT EXISTS automation_records (
|
|
kind TEXT NOT NULL, id TEXT NOT NULL, payload TEXT NOT NULL,
|
|
updated_at REAL NOT NULL, PRIMARY KEY(kind,id));
|
|
CREATE TABLE IF NOT EXISTS automation_jobs (
|
|
session_id TEXT PRIMARY KEY, batch_id TEXT, position INTEGER NOT NULL,
|
|
status TEXT NOT NULL, payload TEXT NOT NULL, updated_at REAL NOT NULL);
|
|
CREATE INDEX IF NOT EXISTS automation_batch_jobs ON automation_jobs(batch_id,position);
|
|
""")
|
|
yield conn
|
|
finally:
|
|
conn.close()
|
|
|
|
def get(self, kind: str, key: str) -> dict | None:
|
|
with self.connection() as conn:
|
|
row = conn.execute(
|
|
"SELECT payload FROM automation_records WHERE kind=? AND id=?", (kind, key)
|
|
).fetchone()
|
|
return json.loads(row[0]) if row else None
|
|
|
|
def put(self, kind: str, key: str, payload: dict) -> dict:
|
|
with self.connection() as conn:
|
|
if kind == "batch":
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
row = conn.execute(
|
|
"SELECT payload FROM automation_records WHERE kind='batch' AND id=?", (key,)
|
|
).fetchone()
|
|
if row is None or json.loads(row[0]).get("status") == "deleting":
|
|
raise ValueError("批次正在删除或已删除，不能继续更新")
|
|
conn.execute(
|
|
"INSERT INTO automation_records VALUES (?,?,?,?) ON CONFLICT(kind,id) DO UPDATE SET payload=excluded.payload,updated_at=excluded.updated_at",
|
|
(kind, key, json.dumps(payload, ensure_ascii=False), time.time()),
|
|
)
|
|
conn.commit()
|
|
return payload
|
|
|
|
def records(self, kind: str, limit: int = 500) -> list[dict]:
|
|
with self.connection() as conn:
|
|
rows = conn.execute(
|
|
"SELECT payload FROM automation_records WHERE kind=? ORDER BY updated_at DESC LIMIT ?",
|
|
(kind, limit),
|
|
).fetchall()
|
|
return [json.loads(row[0]) for row in rows]
|
|
|
|
def job(self, session_id: str) -> dict | None:
|
|
with self.connection() as conn:
|
|
row = conn.execute(
|
|
"SELECT payload,status FROM automation_jobs WHERE session_id=?", (session_id,)
|
|
).fetchone()
|
|
return {**json.loads(row[0]), "status": row[1]} if row else None
|
|
|
|
def save_job(self, payload: dict) -> dict:
|
|
with self.connection() as conn:
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
if batch_id := payload.get("batch_id"):
|
|
row = conn.execute(
|
|
"SELECT payload FROM automation_records WHERE kind='batch' AND id=?",
|
|
(batch_id,),
|
|
).fetchone()
|
|
if row is None or json.loads(row[0]).get("status") == "deleting":
|
|
raise ValueError("批次正在删除，不能继续派发任务")
|
|
conn.execute(
|
|
"INSERT INTO automation_jobs VALUES (?,?,?,?,?,?) ON CONFLICT(session_id) DO UPDATE SET payload=excluded.payload,status=excluded.status,updated_at=excluded.updated_at",
|
|
(
|
|
payload["session_id"],
|
|
payload.get("batch_id"),
|
|
payload.get("position", 0),
|
|
payload["status"],
|
|
json.dumps(payload, ensure_ascii=False),
|
|
time.time(),
|
|
),
|
|
)
|
|
conn.commit()
|
|
return payload
|
|
|
|
def jobs(self, batch_id: str | None = None) -> list[dict]:
|
|
with self.connection() as conn:
|
|
rows = (
|
|
conn.execute(
|
|
"SELECT payload,status FROM automation_jobs WHERE batch_id=? ORDER BY position",
|
|
(batch_id,),
|
|
).fetchall()
|
|
if batch_id
|
|
else conn.execute(
|
|
"SELECT payload,status FROM automation_jobs ORDER BY updated_at DESC LIMIT 10000"
|
|
).fetchall()
|
|
)
|
|
return [{**json.loads(row[0]), "status": row[1]} for row in rows]
|
|
|
|
def jobs_for_sessions(self, session_ids: list[str]) -> dict[str, dict]:
|
|
"""Load only metadata belonging to the sessions shown in history."""
|
|
result = {}
|
|
with self.connection() as conn:
|
|
for offset in range(0, len(session_ids), 400):
|
|
chunk = session_ids[offset : offset + 400]
|
|
markers = ",".join("?" for _ in chunk)
|
|
rows = conn.execute(
|
|
f"SELECT session_id,payload,status FROM automation_jobs WHERE session_id IN ({markers})",
|
|
chunk,
|
|
).fetchall()
|
|
result.update(
|
|
{
|
|
row["session_id"]: {**json.loads(row["payload"]), "status": row["status"]}
|
|
for row in rows
|
|
}
|
|
)
|
|
return result
|
|
|
|
def preferences(self) -> dict:
|
|
return self.get("settings", "reports") or {"formats": ["docx", "xlsx"]}
|
|
|
|
def create_batch(self, batch: dict, jobs: list[dict], payload: dict) -> dict:
|
|
"""Commit a batch and all its jobs atomically; repeated request IDs never duplicate work."""
|
|
request_id = payload.get("request_id")
|
|
fingerprint = hashlib.sha256(
|
|
json.dumps(payload, sort_keys=True, ensure_ascii=False).encode()
|
|
).hexdigest()
|
|
with self.connection() as conn:
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
if request_id:
|
|
identifier(request_id)
|
|
row = conn.execute(
|
|
"SELECT payload FROM automation_records WHERE kind='batch_request' AND id=?",
|
|
(request_id,),
|
|
).fetchone()
|
|
if row:
|
|
previous = json.loads(row[0])
|
|
if previous["fingerprint"] != fingerprint:
|
|
raise ValueError("请求编号已用于不同内容，请使用新的 request_id")
|
|
existing = conn.execute(
|
|
"SELECT payload FROM automation_records WHERE kind='batch' AND id=?",
|
|
(previous["batch_id"],),
|
|
).fetchone()
|
|
if not existing:
|
|
raise ValueError("该请求对应的历史批次已删除，不能自动重新执行")
|
|
return json.loads(existing[0])
|
|
now = time.time()
|
|
conn.execute(
|
|
"INSERT INTO automation_records VALUES (?,?,?,?)",
|
|
("batch", batch["id"], json.dumps(batch, ensure_ascii=False), now),
|
|
)
|
|
for job in jobs:
|
|
conn.execute(
|
|
"INSERT INTO automation_jobs VALUES (?,?,?,?,?,?)",
|
|
(
|
|
job["session_id"],
|
|
batch["id"],
|
|
job["position"],
|
|
job["status"],
|
|
json.dumps(job, ensure_ascii=False),
|
|
now,
|
|
),
|
|
)
|
|
if request_id:
|
|
conn.execute(
|
|
"INSERT INTO automation_records VALUES (?,?,?,?)",
|
|
(
|
|
"batch_request",
|
|
request_id,
|
|
json.dumps({"batch_id": batch["id"], "fingerprint": fingerprint}),
|
|
now,
|
|
),
|
|
)
|
|
conn.commit()
|
|
return batch
|
|
|
|
def begin_delete_batch(self, batch_id: str) -> list[dict] | None:
|
|
"""Reserve an inactive batch for cleanup before any more jobs can start."""
|
|
with self.connection() as conn:
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
row = conn.execute(
|
|
"SELECT payload FROM automation_records WHERE kind='batch' AND id=?",
|
|
(batch_id,),
|
|
).fetchone()
|
|
if row is None:
|
|
return None
|
|
batch = json.loads(row[0])
|
|
active = conn.execute(
|
|
"SELECT 1 FROM automation_jobs WHERE batch_id=? AND status IN "
|
|
"('queued','running','task_paused','resuming','finalizing') LIMIT 1",
|
|
(batch_id,),
|
|
).fetchone()
|
|
if batch["status"] == "running" or active:
|
|
raise ValueError("批次仍有任务执行中或排队中，请等待完成后删除")
|
|
batch["status"] = "deleting"
|
|
conn.execute(
|
|
"UPDATE automation_records SET payload=?,updated_at=? WHERE kind='batch' AND id=?",
|
|
(json.dumps(batch, ensure_ascii=False), time.time(), batch_id),
|
|
)
|
|
jobs = conn.execute(
|
|
"SELECT payload,status FROM automation_jobs WHERE batch_id=? ORDER BY position",
|
|
(batch_id,),
|
|
).fetchall()
|
|
conn.commit()
|
|
return [{**json.loads(job[0]), "status": job[1]} for job in jobs]
|
|
|
|
def delete_batch(self, batch_id: str) -> bool:
|
|
"""Remove a batch record and all of its jobs after cleanup."""
|
|
with self.connection() as conn:
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
row = conn.execute(
|
|
"SELECT payload FROM automation_records WHERE kind='batch' AND id=?",
|
|
(batch_id,),
|
|
).fetchone()
|
|
if row is None:
|
|
return False
|
|
batch = json.loads(row[0])
|
|
if batch.get("status") != "deleting":
|
|
raise ValueError("批次尚未准备删除")
|
|
conn.execute("DELETE FROM automation_jobs WHERE batch_id=?", (batch_id,))
|
|
conn.execute(
|
|
"DELETE FROM automation_records WHERE kind='batch' AND id=?",
|
|
(batch_id,),
|
|
)
|
|
# Keep the request-id tombstone so deletion cannot replay an old batch.
|
|
conn.commit()
|
|
return True
|
|
|
|
def delete_history(self, session_id: str | None = None):
|
|
with self.connection() as conn:
|
|
if session_id:
|
|
conn.execute("DELETE FROM automation_jobs WHERE session_id=?", (session_id,))
|
|
else:
|
|
conn.execute("DELETE FROM automation_jobs")
|
|
conn.execute("DELETE FROM automation_records WHERE kind='batch'")
|
|
conn.commit()
|
|
|
|
|
|
store = AutomationStore()
|