245 lines
8.4 KiB
Python
245 lines
8.4 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import sqlite3
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from core.models import utc_now
|
|
from workflow.config import DEFAULT_WORKFLOW_DB_PATH
|
|
from workflow.models import (
|
|
WorkflowDefinition,
|
|
WorkflowRun,
|
|
WorkflowRunStatus,
|
|
WorkflowStepResult,
|
|
)
|
|
|
|
|
|
class WorkflowStore:
|
|
def __init__(self, db_path: str | Path = DEFAULT_WORKFLOW_DB_PATH) -> None:
|
|
self.db_path = Path(db_path)
|
|
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._ensure_schema()
|
|
|
|
def save_definition(self, definition: WorkflowDefinition) -> None:
|
|
with self._connect() as connection:
|
|
connection.execute(
|
|
"""
|
|
insert into workflow_definitions (id, name, definition_json)
|
|
values (?, ?, ?)
|
|
on conflict(id) do update set
|
|
name = excluded.name,
|
|
definition_json = excluded.definition_json
|
|
""",
|
|
(
|
|
definition.id,
|
|
definition.name,
|
|
json.dumps(definition.to_dict(), ensure_ascii=False),
|
|
),
|
|
)
|
|
|
|
def get_definition(self, definition_id: str) -> WorkflowDefinition | None:
|
|
with self._connect() as connection:
|
|
row = connection.execute(
|
|
"select definition_json from workflow_definitions where id = ?",
|
|
(definition_id,),
|
|
).fetchone()
|
|
if row is None:
|
|
return None
|
|
return WorkflowDefinition.from_dict(json.loads(row["definition_json"]))
|
|
|
|
def create_run(
|
|
self,
|
|
definition_id: str,
|
|
initial_variables: dict[str, Any] | None = None,
|
|
*,
|
|
device_id: str | None = None,
|
|
) -> WorkflowRun:
|
|
definition = self.get_definition(definition_id)
|
|
if definition is None:
|
|
raise KeyError(f"unknown workflow definition {definition_id}")
|
|
run = WorkflowRun(
|
|
definition_id=definition_id,
|
|
status="running",
|
|
current_step_id=definition.entry_step_id,
|
|
variables=dict(initial_variables or {}),
|
|
device_id=device_id,
|
|
)
|
|
with self._connect() as connection:
|
|
connection.execute(
|
|
"""
|
|
insert into workflow_runs (
|
|
id, definition_id, status, current_step_id, device_id,
|
|
variables_json, created_at, updated_at
|
|
) values (?, ?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(
|
|
run.id,
|
|
run.definition_id,
|
|
run.status,
|
|
run.current_step_id,
|
|
run.device_id,
|
|
json.dumps(run.variables, ensure_ascii=False),
|
|
run.created_at.isoformat(),
|
|
run.updated_at.isoformat(),
|
|
),
|
|
)
|
|
return run
|
|
|
|
def append_step_result(
|
|
self,
|
|
run_id: str,
|
|
step_result: WorkflowStepResult,
|
|
) -> None:
|
|
with self._connect() as connection:
|
|
index = (
|
|
connection.execute(
|
|
"select count(*) as count from workflow_step_results where run_id = ?",
|
|
(run_id,),
|
|
).fetchone()["count"]
|
|
+ 1
|
|
)
|
|
connection.execute(
|
|
"""
|
|
insert into workflow_step_results (
|
|
run_id, step_index, step_id, kind, success, detail_json,
|
|
task_id, timestamp
|
|
) values (?, ?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(
|
|
run_id,
|
|
index,
|
|
step_result.step_id,
|
|
step_result.kind,
|
|
1 if step_result.success else 0,
|
|
json.dumps(step_result.detail, ensure_ascii=False),
|
|
step_result.task_id,
|
|
step_result.timestamp.isoformat(),
|
|
),
|
|
)
|
|
|
|
def update_run(
|
|
self,
|
|
run_id: str,
|
|
*,
|
|
status: WorkflowRunStatus | None = None,
|
|
current_step_id: str | None = None,
|
|
variables: dict[str, Any] | None = None,
|
|
) -> None:
|
|
run = self.get_run(run_id)
|
|
if run is None:
|
|
raise KeyError(f"unknown workflow run {run_id}")
|
|
next_status = status or run.status
|
|
next_step_id = current_step_id
|
|
if current_step_id is None and next_status not in {
|
|
"completed",
|
|
"failed",
|
|
"cancelled",
|
|
}:
|
|
next_step_id = run.current_step_id
|
|
next_variables = dict(run.variables if variables is None else variables)
|
|
with self._connect() as connection:
|
|
connection.execute(
|
|
"""
|
|
update workflow_runs
|
|
set status = ?,
|
|
current_step_id = ?,
|
|
variables_json = ?,
|
|
updated_at = ?
|
|
where id = ?
|
|
""",
|
|
(
|
|
next_status,
|
|
next_step_id,
|
|
json.dumps(next_variables, ensure_ascii=False),
|
|
utc_now().isoformat(),
|
|
run_id,
|
|
),
|
|
)
|
|
|
|
def get_run(self, run_id: str) -> WorkflowRun | None:
|
|
with self._connect() as connection:
|
|
run_row = connection.execute(
|
|
"select * from workflow_runs where id = ?",
|
|
(run_id,),
|
|
).fetchone()
|
|
if run_row is None:
|
|
return None
|
|
result_rows = connection.execute(
|
|
"""
|
|
select * from workflow_step_results
|
|
where run_id = ?
|
|
order by step_index
|
|
""",
|
|
(run_id,),
|
|
).fetchall()
|
|
return WorkflowRun.from_dict(
|
|
{
|
|
"id": run_row["id"],
|
|
"definition_id": run_row["definition_id"],
|
|
"status": run_row["status"],
|
|
"current_step_id": run_row["current_step_id"],
|
|
"device_id": run_row["device_id"],
|
|
"variables": json.loads(run_row["variables_json"]),
|
|
"created_at": run_row["created_at"],
|
|
"updated_at": run_row["updated_at"],
|
|
"step_results": [
|
|
{
|
|
"step_id": row["step_id"],
|
|
"kind": row["kind"],
|
|
"success": bool(row["success"]),
|
|
"detail": json.loads(row["detail_json"]),
|
|
"task_id": row["task_id"],
|
|
"timestamp": row["timestamp"],
|
|
}
|
|
for row in result_rows
|
|
],
|
|
}
|
|
)
|
|
|
|
def _ensure_schema(self) -> None:
|
|
with self._connect() as connection:
|
|
connection.execute(
|
|
"""
|
|
create table if not exists workflow_definitions (
|
|
id text primary key,
|
|
name text not null,
|
|
definition_json text not null
|
|
)
|
|
"""
|
|
)
|
|
connection.execute(
|
|
"""
|
|
create table if not exists workflow_runs (
|
|
id text primary key,
|
|
definition_id text not null,
|
|
status text not null,
|
|
current_step_id text,
|
|
device_id text,
|
|
variables_json text not null,
|
|
created_at text not null,
|
|
updated_at text not null
|
|
)
|
|
"""
|
|
)
|
|
connection.execute(
|
|
"""
|
|
create table if not exists workflow_step_results (
|
|
id integer primary key autoincrement,
|
|
run_id text not null,
|
|
step_index integer not null,
|
|
step_id text not null,
|
|
kind text not null,
|
|
success integer not null,
|
|
detail_json text not null,
|
|
task_id text,
|
|
timestamp text not null
|
|
)
|
|
"""
|
|
)
|
|
|
|
def _connect(self) -> sqlite3.Connection:
|
|
connection = sqlite3.connect(self.db_path)
|
|
connection.row_factory = sqlite3.Row
|
|
return connection
|