qwen_agent/task_queue/task_status.py

#!/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()