"""SQLite 数据访问层。 所有持久化逻辑集中在本模块，业务代码只依赖 Database 提供的方法。后续切换 PostgreSQL 时只需替换本层实现，不修改调度器与路由的业务逻辑。 """ from __future__ import annotations import json import sqlite3 from contextlib import contextmanager from pathlib import Path from typing import Any, Iterator from wov_sdk.models import NodeManifest class Database: """SQLite 数据库封装：负责建表以及节点/实例/工作流/任务/产物的 CRUD。""" def __init__(self, path: Path) -> None: """打开数据库并确保父目录存在、表结构已初始化。""" self.path = path self.path.parent.mkdir(parents=True, exist_ok=True) self._init_schema() @contextmanager def _connect(self) -> Iterator[sqlite3.Connection]: """提供带事务提交的数据库连接上下文。""" conn = sqlite3.connect(self.path) # 按列名读取结果，返回 dict 更直观。 conn.row_factory = sqlite3.Row # 开启外键约束，保证子表记录引用有效。 conn.execute("PRAGMA foreign_keys = ON") try: yield conn conn.commit() finally: conn.close() def _init_schema(self) -> None: """创建全部业务表；已存在的表保持不变。""" with self._connect() as conn: conn.executescript( """ -- 节点注册表：保存节点的最新 manifest JSON。 CREATE TABLE IF NOT EXISTS nodes ( id TEXT PRIMARY KEY, name TEXT NOT NULL, version TEXT NOT NULL, capability TEXT NOT NULL, manifest_json TEXT NOT NULL, created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP ); -- 节点实例表：记录 NodeManager 启动的进程及其状态。 CREATE TABLE IF NOT EXISTS node_instances ( id TEXT PRIMARY KEY, node_id TEXT NOT NULL, status TEXT NOT NULL, pid INTEGER, address TEXT, started_at TEXT, last_used_at TEXT, busy_since TEXT, error TEXT, FOREIGN KEY(node_id) REFERENCES nodes(id) ); -- 工作流表：只保存概要信息，完整定义存版本表。 CREATE TABLE IF NOT EXISTS workflows ( id TEXT PRIMARY KEY, name TEXT NOT NULL, description TEXT NOT NULL DEFAULT '', published INTEGER NOT NULL DEFAULT 0, latest_version INTEGER NOT NULL DEFAULT 0, created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP ); -- 工作流版本表：每个版本保存一份 DAG 定义 JSON。 CREATE TABLE IF NOT EXISTS workflow_versions ( id INTEGER PRIMARY KEY AUTOINCREMENT, workflow_id TEXT NOT NULL, version INTEGER NOT NULL, definition_json TEXT NOT NULL, created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, UNIQUE(workflow_id, version), FOREIGN KEY(workflow_id) REFERENCES workflows(id) ); -- 工作流运行表：记录任务从排队到完成/失败的状态机。 CREATE TABLE IF NOT EXISTS workflow_runs ( id TEXT PRIMARY KEY, workflow_id TEXT NOT NULL, workflow_version INTEGER NOT NULL, status TEXT NOT NULL, current_node_id TEXT, progress REAL NOT NULL DEFAULT 0, error TEXT, input_uri TEXT, created_at TEXT NOT NULL, updated_at TEXT NOT NULL, FOREIGN KEY(workflow_id) REFERENCES workflows(id) ); -- 产物表：记录每个任务各节点的输出 URI，按名称唯一。 CREATE TABLE IF NOT EXISTS artifacts ( id INTEGER PRIMARY KEY AUTOINCREMENT, run_id TEXT NOT NULL, node_id TEXT NOT NULL, name TEXT NOT NULL, uri TEXT NOT NULL, mime_type TEXT NOT NULL DEFAULT '', size INTEGER NOT NULL DEFAULT 0, created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, UNIQUE(run_id, name), FOREIGN KEY(run_id) REFERENCES workflow_runs(id) ); """ ) def upsert_node(self, manifest: NodeManifest) -> None: """插入或更新节点注册记录，同名 ID 覆盖为新版本。""" with self._connect() as conn: conn.execute( """ INSERT INTO nodes (id, name, version, capability, manifest_json) VALUES (?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET name = excluded.name, version = excluded.version, capability = excluded.capability, manifest_json = excluded.manifest_json """, ( manifest.id, manifest.name, manifest.version, manifest.capability, json.dumps(manifest.to_dict(), ensure_ascii=False), ), ) def get_node(self, node_id: str) -> NodeManifest | None: """按 ID 读取节点并反序列化为 NodeManifest。""" with self._connect() as conn: row = conn.execute("SELECT manifest_json FROM nodes WHERE id = ?", (node_id,)).fetchone() if row is None: return None return NodeManifest.from_dict(json.loads(row["manifest_json"])) def list_nodes(self) -> list[NodeManifest]: """按 ID 顺序返回全部注册节点。""" with self._connect() as conn: rows = conn.execute("SELECT manifest_json FROM nodes ORDER BY id").fetchall() return [NodeManifest.from_dict(json.loads(row["manifest_json"])) for row in rows] def delete_node(self, node_id: str) -> None: """删除节点及其全部实例记录。""" with self._connect() as conn: # 先删实例再删节点，满足外键约束。 conn.execute("DELETE FROM node_instances WHERE node_id = ?", (node_id,)) conn.execute("DELETE FROM nodes WHERE id = ?", (node_id,)) def upsert_instance(self, instance: dict[str, Any]) -> None: """插入或更新节点实例状态记录。""" with self._connect() as conn: conn.execute( """ INSERT INTO node_instances ( id, node_id, status, pid, address, started_at, last_used_at, busy_since, error ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET status = excluded.status, pid = excluded.pid, address = excluded.address, started_at = excluded.started_at, last_used_at = excluded.last_used_at, busy_since = excluded.busy_since, error = excluded.error """, ( instance["id"], instance["node_id"], instance["status"], instance.get("pid"), instance.get("address"), instance.get("started_at"), instance.get("last_used_at"), instance.get("busy_since"), instance.get("error"), ), ) def list_instances(self) -> list[dict[str, Any]]: """按启动时间倒序返回全部节点实例。""" with self._connect() as conn: rows = conn.execute( "SELECT * FROM node_instances ORDER BY started_at DESC" ).fetchall() return [dict(row) for row in rows] def delete_instance(self, instance_id: str) -> None: """删除指定实例记录。""" with self._connect() as conn: conn.execute("DELETE FROM node_instances WHERE id = ?", (instance_id,)) def upsert_workflow(self, workflow: dict[str, Any]) -> None: """插入或更新工作流概要信息。""" with self._connect() as conn: conn.execute( """ INSERT INTO workflows (id, name, description, published, latest_version) VALUES (?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET name = excluded.name, description = excluded.description, published = excluded.published, latest_version = excluded.latest_version """, ( workflow["id"], workflow["name"], workflow.get("description", ""), int(workflow.get("published", 0)), int(workflow.get("latest_version", 0)), ), ) def get_workflow(self, workflow_id: str) -> dict[str, Any] | None: """按 ID 读取工作流概要。""" with self._connect() as conn: row = conn.execute("SELECT * FROM workflows WHERE id = ?", (workflow_id,)).fetchone() return dict(row) if row else None def list_workflows(self) -> list[dict[str, Any]]: """按创建时间倒序返回全部工作流。""" with self._connect() as conn: rows = conn.execute("SELECT * FROM workflows ORDER BY created_at DESC").fetchall() return [dict(row) for row in rows] def delete_workflow(self, workflow_id: str) -> None: """级联删除工作流相关的产物、任务、版本和概要记录。""" with self._connect() as conn: # 外键没有级联删除配置，手动按依赖顺序清理。 conn.execute("DELETE FROM artifacts WHERE run_id IN (SELECT id FROM workflow_runs WHERE workflow_id = ?)", (workflow_id,)) conn.execute("DELETE FROM workflow_runs WHERE workflow_id = ?", (workflow_id,)) conn.execute("DELETE FROM workflow_versions WHERE workflow_id = ?", (workflow_id,)) conn.execute("DELETE FROM workflows WHERE id = ?", (workflow_id,)) def create_workflow_version(self, workflow_id: str, version: int, definition: dict[str, Any]) -> None: """为工作流新增一个版本，definition 以 JSON 保存。""" with self._connect() as conn: conn.execute( """ INSERT INTO workflow_versions (workflow_id, version, definition_json) VALUES (?, ?, ?) """, (workflow_id, version, json.dumps(definition, ensure_ascii=False)), ) def get_latest_workflow_version(self, workflow_id: str) -> dict[str, Any] | None: """返回工作流最新版本，并把 definition_json 反序列化为 definition。""" with self._connect() as conn: row = conn.execute( """ SELECT * FROM workflow_versions WHERE workflow_id = ? ORDER BY version DESC LIMIT 1 """, (workflow_id,), ).fetchone() if row is None: return None result = dict(row) # 对外统一暴露 definition 字典，隐藏 JSON 存储细节。 result["definition"] = json.loads(result.pop("definition_json")) return result def get_workflow_version(self, workflow_id: str, version: int) -> dict[str, Any] | None: """按版本号读取指定工作流版本。""" with self._connect() as conn: row = conn.execute( """ SELECT * FROM workflow_versions WHERE workflow_id = ? AND version = ? """, (workflow_id, version), ).fetchone() if row is None: return None result = dict(row) result["definition"] = json.loads(result.pop("definition_json")) return result def list_workflow_versions(self, workflow_id: str) -> list[dict[str, Any]]: """按版本倒序返回工作流全部版本。""" with self._connect() as conn: rows = conn.execute( """ SELECT * FROM workflow_versions WHERE workflow_id = ? ORDER BY version DESC """, (workflow_id,), ).fetchall() versions = [] for row in rows: item = dict(row) item["definition"] = json.loads(item.pop("definition_json")) versions.append(item) return versions def create_run(self, run: dict[str, Any]) -> None: """创建一条排队中的工作流运行记录。""" with self._connect() as conn: conn.execute( """ INSERT INTO workflow_runs ( id, workflow_id, workflow_version, status, current_node_id, progress, error, input_uri, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( run["id"], run["workflow_id"], run["workflow_version"], run["status"], run.get("current_node_id"), float(run.get("progress", 0)), run.get("error"), run.get("input_uri"), run["created_at"], run["updated_at"], ), ) def get_run(self, run_id: str) -> dict[str, Any] | None: """按 ID 读取任务运行记录。""" with self._connect() as conn: row = conn.execute("SELECT * FROM workflow_runs WHERE id = ?", (run_id,)).fetchone() return dict(row) if row else None def list_runs(self, limit: int = 20) -> list[dict[str, Any]]: """按创建时间倒序返回最近的运行记录。""" with self._connect() as conn: rows = conn.execute( "SELECT * FROM workflow_runs ORDER BY created_at DESC LIMIT ?", (limit,), ).fetchall() return [dict(row) for row in rows] def update_run(self, run_id: str, **fields: Any) -> None: """更新运行状态字段，同时刷新 updated_at；未知字段会被忽略。""" # 只允许更新状态机相关字段，防止任意列被改写。 allowed = { "status", "current_node_id", "progress", "error", } updates = {key: value for key, value in fields.items() if key in allowed} if not updates: return updates["updated_at"] = fields.get("updated_at") # 动态拼接 SET 子句，键来自白名单，不存在 SQL 注入风险。 assignments = ", ".join(f"{key} = ?" for key in updates) values = list(updates.values()) + [run_id] with self._connect() as conn: conn.execute(f"UPDATE workflow_runs SET {assignments} WHERE id = ?", values) def reset_run(self, run_id: str, updated_at: str) -> None: """把失败任务重置为 QUEUED，并清空进度与旧产物，供重试使用。""" with self._connect() as conn: # 清空错误和进度，恢复到首次排队时的状态。 conn.execute( """ UPDATE workflow_runs SET status = 'QUEUED', current_node_id = NULL, progress = 0, error = NULL, updated_at = ? WHERE id = ? """, (updated_at, run_id), ) # 删除旧产物，避免重试后残留过期下载链接。 conn.execute("DELETE FROM artifacts WHERE run_id = ?", (run_id,)) def next_queued_run(self) -> dict[str, Any] | None: """按创建时间返回最早一条排队任务，供调度器轮询。""" with self._connect() as conn: row = conn.execute( """ SELECT * FROM workflow_runs WHERE status = 'QUEUED' ORDER BY created_at ASC LIMIT 1 """ ).fetchone() return dict(row) if row else None def create_artifact(self, artifact: dict[str, Any]) -> None: """记录任务产物；同 run 与 name 冲突时覆盖。""" with self._connect() as conn: conn.execute( """ INSERT OR REPLACE INTO artifacts ( run_id, node_id, name, uri, mime_type, size ) VALUES (?, ?, ?, ?, ?, ?) """, ( artifact["run_id"], artifact["node_id"], artifact["name"], artifact["uri"], artifact.get("mime_type", ""), int(artifact.get("size", 0)), ), ) def list_artifacts(self, run_id: str) -> list[dict[str, Any]]: """按创建时间返回任务的全部产物。""" with self._connect() as conn: rows = conn.execute( "SELECT * FROM artifacts WHERE run_id = ? ORDER BY created_at", (run_id,), ).fetchall() return [dict(row) for row in rows] def get_artifact(self, run_id: str, name: str) -> dict[str, Any] | None: """按任务与产物名读取单个产物记录。""" with self._connect() as conn: row = conn.execute( "SELECT * FROM artifacts WHERE run_id = ? AND name = ?", (run_id, name), ).fetchone() return dict(row) if row else None def delete_run_artifacts(self, run_id: str) -> None: """删除任务的全部产物记录。""" with self._connect() as conn: conn.execute("DELETE FROM artifacts WHERE run_id = ?", (run_id,))