158 lines
6.7 KiB
Python
158 lines
6.7 KiB
Python
# [DEF:TaskPersistenceModule:Module]
|
|
# @SEMANTICS: persistence, sqlite, task, storage
|
|
# @PURPOSE: Handles the persistence of tasks, specifically those awaiting user input, to a SQLite database.
|
|
# @LAYER: Core
|
|
# @RELATION: Used by TaskManager to save and load tasks.
|
|
# @INVARIANT: Database schema must match the Task model structure.
|
|
# @CONSTRAINT: Uses synchronous SQLite operations (blocking), should be used carefully.
|
|
|
|
# [SECTION: IMPORTS]
|
|
import sqlite3
|
|
import json
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Any
|
|
|
|
from .models import Task, TaskStatus
|
|
from ..logger import logger, belief_scope
|
|
# [/SECTION]
|
|
|
|
# [DEF:TaskPersistenceService:Class]
|
|
# @SEMANTICS: persistence, service, database
|
|
# @PURPOSE: Provides methods to save and load tasks from a local SQLite database.
|
|
class TaskPersistenceService:
|
|
def __init__(self, db_path: Optional[Path] = None):
|
|
if db_path is None:
|
|
self.db_path = Path(__file__).parent.parent.parent.parent / "migrations.db"
|
|
else:
|
|
self.db_path = db_path
|
|
self._ensure_db_exists()
|
|
|
|
# [DEF:TaskPersistenceService._ensure_db_exists:Function]
|
|
# @PURPOSE: Ensures the database directory and table exist.
|
|
# @PRE: None.
|
|
# @POST: Database file and table are created if they didn't exist.
|
|
def _ensure_db_exists(self) -> None:
|
|
with belief_scope("TaskPersistenceService._ensure_db_exists"):
|
|
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS persistent_tasks (
|
|
id TEXT PRIMARY KEY,
|
|
plugin_id TEXT NOT NULL,
|
|
status TEXT NOT NULL,
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL,
|
|
input_request JSON,
|
|
context JSON
|
|
)
|
|
""")
|
|
conn.commit()
|
|
conn.close()
|
|
# [/DEF:TaskPersistenceService._ensure_db_exists:Function]
|
|
|
|
# [DEF:TaskPersistenceService.persist_tasks:Function]
|
|
# @PURPOSE: Persists a list of tasks to the database.
|
|
# @PRE: Tasks list contains valid Task objects.
|
|
# @POST: Tasks matching the criteria (AWAITING_INPUT) are saved/updated in the DB.
|
|
# @PARAM: tasks (List[Task]) - The list of tasks to check and persist.
|
|
def persist_tasks(self, tasks: List[Task]) -> None:
|
|
with belief_scope("TaskPersistenceService.persist_tasks"):
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
count = 0
|
|
for task in tasks:
|
|
if task.status == TaskStatus.AWAITING_INPUT:
|
|
cursor.execute("""
|
|
INSERT OR REPLACE INTO persistent_tasks
|
|
(id, plugin_id, status, created_at, updated_at, input_request, context)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?)
|
|
""", (
|
|
task.id,
|
|
task.plugin_id,
|
|
task.status.value,
|
|
task.started_at.isoformat() if task.started_at else datetime.utcnow().isoformat(),
|
|
datetime.utcnow().isoformat(),
|
|
json.dumps(task.input_request) if task.input_request else None,
|
|
json.dumps(task.params)
|
|
))
|
|
count += 1
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
logger.info(f"Persisted {count} tasks awaiting input.")
|
|
# [/DEF:TaskPersistenceService.persist_tasks:Function]
|
|
|
|
# [DEF:TaskPersistenceService.load_tasks:Function]
|
|
# @PURPOSE: Loads persisted tasks from the database.
|
|
# @PRE: Database exists.
|
|
# @POST: Returns a list of Task objects reconstructed from the DB.
|
|
# @RETURN: List[Task] - The loaded tasks.
|
|
def load_tasks(self) -> List[Task]:
|
|
with belief_scope("TaskPersistenceService.load_tasks"):
|
|
if not self.db_path.exists():
|
|
return []
|
|
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
|
|
# Check if plugin_id column exists (migration for existing db)
|
|
cursor.execute("PRAGMA table_info(persistent_tasks)")
|
|
columns = [info[1] for info in cursor.fetchall()]
|
|
has_plugin_id = "plugin_id" in columns
|
|
|
|
if has_plugin_id:
|
|
cursor.execute("SELECT id, plugin_id, status, created_at, input_request, context FROM persistent_tasks")
|
|
else:
|
|
cursor.execute("SELECT id, status, created_at, input_request, context FROM persistent_tasks")
|
|
|
|
rows = cursor.fetchall()
|
|
|
|
loaded_tasks = []
|
|
for row in rows:
|
|
if has_plugin_id:
|
|
task_id, plugin_id, status, created_at, input_request_json, context_json = row
|
|
else:
|
|
task_id, status, created_at, input_request_json, context_json = row
|
|
plugin_id = "superset-migration" # Default fallback
|
|
|
|
try:
|
|
task = Task(
|
|
id=task_id,
|
|
plugin_id=plugin_id,
|
|
status=TaskStatus(status),
|
|
started_at=datetime.fromisoformat(created_at),
|
|
input_required=True,
|
|
input_request=json.loads(input_request_json) if input_request_json else None,
|
|
params=json.loads(context_json) if context_json else {}
|
|
)
|
|
loaded_tasks.append(task)
|
|
except Exception as e:
|
|
logger.error(f"Failed to load task {task_id}: {e}")
|
|
|
|
conn.close()
|
|
return loaded_tasks
|
|
# [/DEF:TaskPersistenceService.load_tasks:Function]
|
|
|
|
# [DEF:TaskPersistenceService.delete_tasks:Function]
|
|
# @PURPOSE: Deletes specific tasks from the database.
|
|
# @PARAM: task_ids (List[str]) - List of task IDs to delete.
|
|
def delete_tasks(self, task_ids: List[str]) -> None:
|
|
if not task_ids:
|
|
return
|
|
with belief_scope("TaskPersistenceService.delete_tasks"):
|
|
conn = sqlite3.connect(str(self.db_path))
|
|
cursor = conn.cursor()
|
|
placeholders = ', '.join('?' for _ in task_ids)
|
|
cursor.execute(f"DELETE FROM persistent_tasks WHERE id IN ({placeholders})", task_ids)
|
|
conn.commit()
|
|
conn.close()
|
|
# [/DEF:TaskPersistenceService.delete_tasks:Function]
|
|
|
|
# [/DEF:TaskPersistenceService:Class]
|
|
|
|
# [/DEF:TaskPersistenceModule:Module] |