158 lines
6.7 KiB
Python
158 lines
6.7 KiB
Python
# [DEF:TaskPersistenceModule:Module]
|
|
# @SEMANTICS: persistence, sqlite, sqlalchemy, task, storage
|
|
# @PURPOSE: Handles the persistence of tasks using SQLAlchemy and the tasks.db database.
|
|
# @LAYER: Core
|
|
# @RELATION: Used by TaskManager to save and load tasks.
|
|
# @INVARIANT: Database schema must match the TaskRecord model structure.
|
|
|
|
# [SECTION: IMPORTS]
|
|
from datetime import datetime
|
|
from typing import List, Optional, Dict, Any
|
|
import json
|
|
|
|
from sqlalchemy.orm import Session
|
|
from ...models.task import TaskRecord
|
|
from ..database import TasksSessionLocal
|
|
from .models import Task, TaskStatus, LogEntry
|
|
from ..logger import logger, belief_scope
|
|
# [/SECTION]
|
|
|
|
# [DEF:TaskPersistenceService:Class]
|
|
# @SEMANTICS: persistence, service, database, sqlalchemy
|
|
# @PURPOSE: Provides methods to save and load tasks from the tasks.db database using SQLAlchemy.
|
|
class TaskPersistenceService:
|
|
# [DEF:__init__:Function]
|
|
# @PURPOSE: Initializes the persistence service.
|
|
# @PRE: None.
|
|
# @POST: Service is ready.
|
|
def __init__(self):
|
|
with belief_scope("TaskPersistenceService.__init__"):
|
|
# We use TasksSessionLocal from database.py
|
|
pass
|
|
# [/DEF:__init__:Function]
|
|
|
|
# [DEF:persist_task:Function]
|
|
# @PURPOSE: Persists or updates a single task in the database.
|
|
# @PRE: isinstance(task, Task)
|
|
# @POST: Task record created or updated in database.
|
|
# @PARAM: task (Task) - The task object to persist.
|
|
def persist_task(self, task: Task) -> None:
|
|
with belief_scope("TaskPersistenceService.persist_task", f"task_id={task.id}"):
|
|
session: Session = TasksSessionLocal()
|
|
try:
|
|
record = session.query(TaskRecord).filter(TaskRecord.id == task.id).first()
|
|
if not record:
|
|
record = TaskRecord(id=task.id)
|
|
session.add(record)
|
|
|
|
record.type = task.plugin_id
|
|
record.status = task.status.value
|
|
record.environment_id = task.params.get("environment_id") or task.params.get("source_env_id")
|
|
record.started_at = task.started_at
|
|
record.finished_at = task.finished_at
|
|
record.params = task.params
|
|
record.result = task.result
|
|
|
|
# Store logs as JSON, converting datetime to string
|
|
record.logs = []
|
|
for log in task.logs:
|
|
log_dict = log.dict()
|
|
if isinstance(log_dict.get('timestamp'), datetime):
|
|
log_dict['timestamp'] = log_dict['timestamp'].isoformat()
|
|
record.logs.append(log_dict)
|
|
|
|
# Extract error if failed
|
|
if task.status == TaskStatus.FAILED:
|
|
for log in reversed(task.logs):
|
|
if log.level == "ERROR":
|
|
record.error = log.message
|
|
break
|
|
|
|
session.commit()
|
|
except Exception as e:
|
|
session.rollback()
|
|
logger.error(f"Failed to persist task {task.id}: {e}")
|
|
finally:
|
|
session.close()
|
|
# [/DEF:persist_task:Function]
|
|
|
|
# [DEF:persist_tasks:Function]
|
|
# @PURPOSE: Persists multiple tasks.
|
|
# @PRE: isinstance(tasks, list)
|
|
# @POST: All tasks in list are persisted.
|
|
# @PARAM: tasks (List[Task]) - The list of tasks to persist.
|
|
def persist_tasks(self, tasks: List[Task]) -> None:
|
|
with belief_scope("TaskPersistenceService.persist_tasks"):
|
|
for task in tasks:
|
|
self.persist_task(task)
|
|
# [/DEF:persist_tasks:Function]
|
|
|
|
# [DEF:load_tasks:Function]
|
|
# @PURPOSE: Loads tasks from the database.
|
|
# @PRE: limit is an integer.
|
|
# @POST: Returns list of Task objects.
|
|
# @PARAM: limit (int) - Max tasks to load.
|
|
# @PARAM: status (Optional[TaskStatus]) - Filter by status.
|
|
# @RETURN: List[Task] - The loaded tasks.
|
|
def load_tasks(self, limit: int = 100, status: Optional[TaskStatus] = None) -> List[Task]:
|
|
with belief_scope("TaskPersistenceService.load_tasks"):
|
|
session: Session = TasksSessionLocal()
|
|
try:
|
|
query = session.query(TaskRecord)
|
|
if status:
|
|
query = query.filter(TaskRecord.status == status.value)
|
|
|
|
records = query.order_by(TaskRecord.created_at.desc()).limit(limit).all()
|
|
|
|
loaded_tasks = []
|
|
for record in records:
|
|
try:
|
|
logs = []
|
|
if record.logs:
|
|
for log_data in record.logs:
|
|
# Handle timestamp conversion if it's a string
|
|
if isinstance(log_data.get('timestamp'), str):
|
|
log_data['timestamp'] = datetime.fromisoformat(log_data['timestamp'])
|
|
logs.append(LogEntry(**log_data))
|
|
|
|
task = Task(
|
|
id=record.id,
|
|
plugin_id=record.type,
|
|
status=TaskStatus(record.status),
|
|
started_at=record.started_at,
|
|
finished_at=record.finished_at,
|
|
params=record.params or {},
|
|
result=record.result,
|
|
logs=logs
|
|
)
|
|
loaded_tasks.append(task)
|
|
except Exception as e:
|
|
logger.error(f"Failed to reconstruct task {record.id}: {e}")
|
|
|
|
return loaded_tasks
|
|
finally:
|
|
session.close()
|
|
# [/DEF:load_tasks:Function]
|
|
|
|
# [DEF:delete_tasks:Function]
|
|
# @PURPOSE: Deletes specific tasks from the database.
|
|
# @PRE: task_ids is a list of strings.
|
|
# @POST: Specified task records deleted from 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"):
|
|
session: Session = TasksSessionLocal()
|
|
try:
|
|
session.query(TaskRecord).filter(TaskRecord.id.in_(task_ids)).delete(synchronize_session=False)
|
|
session.commit()
|
|
except Exception as e:
|
|
session.rollback()
|
|
logger.error(f"Failed to delete tasks: {e}")
|
|
finally:
|
|
session.close()
|
|
# [/DEF:delete_tasks:Function]
|
|
|
|
# [/DEF:TaskPersistenceService:Class]
|
|
# [/DEF:TaskPersistenceModule:Module] |