#!/usr/bin/env python3 """ Task status SQLite storage system. """ import json import os import sqlite3 import time from typing import Dict, Optional, Any, List from pathlib import Path class TaskStatusStore: """SQLite-based task status store.""" def __init__(self, db_path: str = "projects/queue_data/task_status.db"): self.db_path = db_path # Ensure directory exists Path(db_path).parent.mkdir(parents=True, exist_ok=True) self._init_database() def _init_database(self): """Initialize database tables.""" with sqlite3.connect(self.db_path) as conn: conn.execute(''' CREATE TABLE IF NOT EXISTS task_status ( task_id TEXT PRIMARY KEY, unique_id TEXT NOT NULL, status TEXT NOT NULL, created_at REAL NOT NULL, updated_at REAL NOT NULL, result TEXT, error TEXT ) ''') conn.commit() def set_status(self, task_id: str, unique_id: str, status: str, result: Optional[Dict] = None, error: Optional[str] = None): """Set task status.""" current_time = time.time() with sqlite3.connect(self.db_path) as conn: conn.execute(''' INSERT OR REPLACE INTO task_status (task_id, unique_id, status, created_at, updated_at, result, error) VALUES (?, ?, ?, ?, ?, ?, ?) ''', ( task_id, unique_id, status, current_time, current_time, json.dumps(result) if result else None, error )) conn.commit() def get_status(self, task_id: str) -> Optional[Dict]: """Get task status.""" with sqlite3.connect(self.db_path) as conn: conn.row_factory = sqlite3.Row cursor = conn.execute( 'SELECT * FROM task_status WHERE task_id = ?', (task_id,) ) row = cursor.fetchone() if not row: return None result = dict(row) # Parse JSON field if result['result']: result['result'] = json.loads(result['result']) return result def update_status(self, task_id: str, status: str, result: Optional[Dict] = None, error: Optional[str] = None): """Update task status.""" with sqlite3.connect(self.db_path) as conn: # Check if task exists cursor = conn.execute( 'SELECT task_id FROM task_status WHERE task_id = ?', (task_id,) ) if not cursor.fetchone(): return False # Update status conn.execute(''' UPDATE task_status SET status = ?, updated_at = ?, result = ?, error = ? WHERE task_id = ? ''', ( status, time.time(), json.dumps(result) if result else None, error, task_id )) conn.commit() return True def delete_status(self, task_id: str): """Delete task status.""" with sqlite3.connect(self.db_path) as conn: cursor = conn.execute( 'DELETE FROM task_status WHERE task_id = ?', (task_id,) ) conn.commit() return cursor.rowcount > 0 def list_all(self) -> Dict[str, Dict]: """List all task statuses.""" with sqlite3.connect(self.db_path) as conn: conn.row_factory = sqlite3.Row cursor = conn.execute( 'SELECT * FROM task_status ORDER BY updated_at DESC' ) all_tasks = {} for row in cursor: result = dict(row) # Parse JSON field if result['result']: result['result'] = json.loads(result['result']) all_tasks[result['task_id']] = result return all_tasks def get_by_unique_id(self, unique_id: str) -> List[Dict]: """Get all tasks for a given project ID.""" with sqlite3.connect(self.db_path) as conn: conn.row_factory = sqlite3.Row cursor = conn.execute( 'SELECT * FROM task_status WHERE unique_id = ? ORDER BY updated_at DESC', (unique_id,) ) tasks = [] for row in cursor: result = dict(row) if result['result']: result['result'] = json.loads(result['result']) tasks.append(result) return tasks def cleanup_old_tasks(self, older_than_days: int = 7) -> int: """Clean up old task records.""" cutoff_time = time.time() - (older_than_days * 24 * 3600) with sqlite3.connect(self.db_path) as conn: cursor = conn.execute( 'DELETE FROM task_status WHERE updated_at < ?', (cutoff_time,) ) conn.commit() return cursor.rowcount def get_statistics(self) -> Dict[str, Any]: """Get task statistics.""" with sqlite3.connect(self.db_path) as conn: # Total tasks total = conn.execute('SELECT COUNT(*) FROM task_status').fetchone()[0] # Status breakdown status_stats = conn.execute(''' SELECT status, COUNT(*) as count FROM task_status GROUP BY status ''').fetchall() # Tasks in the last 24 hours recent = time.time() - (24 * 3600) recent_tasks = conn.execute( 'SELECT COUNT(*) FROM task_status WHERE updated_at > ?', (recent,) ).fetchone()[0] return { 'total_tasks': total, 'status_breakdown': dict(status_stats), 'recent_24h': recent_tasks, 'database_path': self.db_path } def search_tasks(self, status: Optional[str] = None, unique_id: Optional[str] = None, limit: int = 100) -> List[Dict]: """Search tasks.""" query = 'SELECT * FROM task_status WHERE 1=1' params = [] if status: query += ' AND status = ?' params.append(status) if unique_id: query += ' AND unique_id = ?' params.append(unique_id) query += ' ORDER BY updated_at DESC LIMIT ?' params.append(limit) with sqlite3.connect(self.db_path) as conn: conn.row_factory = sqlite3.Row cursor = conn.execute(query, params) tasks = [] for row in cursor: result = dict(row) if result['result']: result['result'] = json.loads(result['result']) tasks.append(result) return tasks # Global task status store instance task_status_store = TaskStatusStore()