from __future__ import annotations import json import sqlite3 import uuid from datetime import datetime from pathlib import Path from typing import Any DB_PATH = Path(__file__).resolve().parents[1] / 'task_db.sqlite' class TaskManager: def __init__(self, db_path: str | None = None) -> None: self.db_path = db_path or str(DB_PATH) self._init_db() def _conn(self): return sqlite3.connect(self.db_path) def _init_db(self) -> None: with self._conn() as c: c.execute(''' CREATE TABLE IF NOT EXISTS tasks ( task_id TEXT PRIMARY KEY, status TEXT NOT NULL, payload_json TEXT, result_json TEXT, created_at TEXT, updated_at TEXT ) ''') def create_task(self, payload: dict[str, Any]) -> str: task_id = uuid.uuid4().hex now = datetime.now().strftime('%Y-%m-%d %H:%M:%S') with self._conn() as c: c.execute( 'INSERT INTO tasks(task_id,status,payload_json,result_json,created_at,updated_at) VALUES(?,?,?,?,?,?)', (task_id, 'queued', json.dumps(payload, ensure_ascii=False), '{}', now, now), ) return task_id def cancel_task(self, task_id: str) -> bool: now = datetime.now().strftime('%Y-%m-%d %H:%M:%S') with self._conn() as c: cur = c.execute( "UPDATE tasks SET status='cancelled',updated_at=? WHERE task_id=? AND status IN ('queued','running')", (now, task_id), ) return cur.rowcount > 0 def get_task(self, task_id: str) -> dict[str, Any] | None: with self._conn() as c: row = c.execute( 'SELECT task_id,status,payload_json,result_json,created_at,updated_at FROM tasks WHERE task_id=?', (task_id,), ).fetchone() if not row: return None return { 'task_id': row[0], 'status': row[1], 'payload': json.loads(row[2] or '{}'), 'result': json.loads(row[3] or '{}'), 'created_at': row[4], 'updated_at': row[5], } def list_tasks(self, limit: int = 100) -> list[dict[str, Any]]: with self._conn() as c: rows = c.execute( 'SELECT task_id,status,payload_json,result_json,created_at,updated_at FROM tasks ORDER BY created_at DESC LIMIT ?', (limit,), ).fetchall() return [ { 'task_id': r[0], 'status': r[1], 'payload': json.loads(r[2] or '{}'), 'result': json.loads(r[3] or '{}'), 'created_at': r[4], 'updated_at': r[5], } for r in rows ]