396 lines
15 KiB
Python
396 lines
15 KiB
Python
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:
|
|
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)
|
|
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(
|
|
"""
|
|
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
|
|
);
|
|
|
|
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
|
|
);
|
|
|
|
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)
|
|
);
|
|
|
|
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:
|
|
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:
|
|
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]:
|
|
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:
|
|
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:
|
|
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:
|
|
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)
|
|
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:
|
|
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:
|
|
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")
|
|
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:
|
|
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:
|
|
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,))
|