agyqueue 0.1.0__py3-none-any.whl

agyqueue/storage.py

@@ -0,0 +1,423 @@
+import sqlite3
+import os
+import logging
+from abc import ABC, abstractmethod
+from datetime import datetime, timezone
+from typing import Optional, List
+from contextlib import contextmanager
+from agyqueue.models import Task, TaskStatus
+
+logger = logging.getLogger(__name__)
+
+class BaseTaskStore(ABC):
+    """Abstract base class defining the interface for task storage backends."""
+    
+    @abstractmethod
+    def save_task(self, task: Task) -> None:
+        """Saves a new task or updates an existing one."""
+        pass
+
+    @abstractmethod
+    def get_task(self, task_id: str) -> Optional[Task]:
+        """Retrieves a task by its unique ID."""
+        pass
+
+    @abstractmethod
+    def update_task(
+        self,
+        task_id: str,
+        status: TaskStatus,
+        progress: int,
+        step: str,
+        result: Optional[str] = None,
+        error: Optional[str] = None,
+    ) -> Optional[Task]:
+        """Updates the status, progress, current step, and optionally result/error of a task."""
+        pass
+
+    @abstractmethod
+    def touch_task(self, task_id: str) -> None:
+        """Touches the task to update its updated_at heartbeat timestamp."""
+        pass
+
+    @abstractmethod
+    def list_tasks(self) -> List[Task]:
+        """Lists all tasks ordered by creation date desc."""
+        pass
+
+    @abstractmethod
+    def get_subtasks(self, parent_id: str) -> List[Task]:
+        """Lists all child subtasks for a given parent task."""
+        pass
+
+
+class SQLiteTaskStore(BaseTaskStore):
+    """SQLite implementation of the TaskStore interface."""
+    
+    def __init__(self, db_path: Optional[str] = None):
+        self.db_path = db_path or os.environ.get("AGYQUEUE_DB_PATH", "agyqueue.db")
+        self._init_db()
+
+    def _get_conn(self) -> sqlite3.Connection:
+        conn = sqlite3.connect(self.db_path, timeout=30.0)
+        # Enable WAL mode for better concurrency in multi-process environments
+        conn.execute("PRAGMA journal_mode=WAL;")
+        conn.row_factory = sqlite3.Row
+        return conn
+
+    def _init_db(self):
+        with self._get_conn() as conn:
+            conn.execute("""
+                CREATE TABLE IF NOT EXISTS tasks (
+                    task_id TEXT PRIMARY KEY,
+                    prompt TEXT NOT NULL,
+                    task_type TEXT NOT NULL,
+                    status TEXT NOT NULL,
+                    progress INTEGER NOT NULL,
+                    step TEXT NOT NULL,
+                    result TEXT,
+                    error TEXT,
+                    parent_id TEXT,
+                    created_at TEXT NOT NULL,
+                    updated_at TEXT NOT NULL
+                )
+            """)
+            conn.commit()
+            
+            # Schema migration: add parent_id if table already existed without it
+            try:
+                conn.execute("ALTER TABLE tasks ADD COLUMN parent_id TEXT;")
+                conn.commit()
+            except sqlite3.OperationalError:
+                # Column already exists
+                pass
+
+    def save_task(self, task: Task) -> None:
+        with self._get_conn() as conn:
+            conn.execute(
+                """
+                INSERT OR REPLACE INTO tasks 
+                (task_id, prompt, task_type, status, progress, step, result, error, parent_id, created_at, updated_at)
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                """,
+                (
+                    task.task_id,
+                    task.prompt,
+                    task.task_type,
+                    task.status.value,
+                    task.progress,
+                    task.step,
+                    task.result,
+                    task.error,
+                    task.parent_id,
+                    task.created_at,
+                    task.updated_at,
+                ),
+            )
+            conn.commit()
+
+    def get_task(self, task_id: str) -> Optional[Task]:
+        with self._get_conn() as conn:
+            row = conn.execute(
+                "SELECT * FROM tasks WHERE task_id = ?", (task_id,)
+            ).fetchone()
+            if not row:
+                return None
+            return Task(
+                task_id=row["task_id"],
+                prompt=row["prompt"],
+                task_type=row["task_type"],
+                status=TaskStatus(row["status"]),
+                progress=row["progress"],
+                step=row["step"],
+                result=row["result"],
+                error=row["error"],
+                parent_id=row["parent_id"],
+                created_at=row["created_at"],
+                updated_at=row["updated_at"],
+            )
+
+    def update_task(
+        self,
+        task_id: str,
+        status: TaskStatus,
+        progress: int,
+        step: str,
+        result: Optional[str] = None,
+        error: Optional[str] = None,
+    ) -> Optional[Task]:
+        now = datetime.now(timezone.utc).replace(tzinfo=None).isoformat()
+        with self._get_conn() as conn:
+            conn.execute(
+                """
+                UPDATE tasks
+                SET status = ?, progress = ?, step = ?, result = COALESCE(?, result), error = COALESCE(?, error), updated_at = ?
+                WHERE task_id = ?
+                """,
+                (status.value, progress, step, result, error, now, task_id),
+            )
+            conn.commit()
+        return self.get_task(task_id)
+
+    def touch_task(self, task_id: str) -> None:
+        now = datetime.now(timezone.utc).replace(tzinfo=None).isoformat()
+        with self._get_conn() as conn:
+            conn.execute(
+                "UPDATE tasks SET updated_at = ? WHERE task_id = ?",
+                (now, task_id),
+            )
+            conn.commit()
+
+    def list_tasks(self) -> List[Task]:
+        with self._get_conn() as conn:
+            rows = conn.execute("SELECT * FROM tasks ORDER BY created_at DESC").fetchall()
+            return [
+                Task(
+                    task_id=row["task_id"],
+                    prompt=row["prompt"],
+                    task_type=row["task_type"],
+                    status=TaskStatus(row["status"]),
+                    progress=row["progress"],
+                    step=row["step"],
+                    result=row["result"],
+                    error=row["error"],
+                    parent_id=row["parent_id"],
+                    created_at=row["created_at"],
+                    updated_at=row["updated_at"],
+                )
+                for row in rows
+            ]
+
+    def get_subtasks(self, parent_id: str) -> List[Task]:
+        with self._get_conn() as conn:
+            rows = conn.execute(
+                "SELECT * FROM tasks WHERE parent_id = ? ORDER BY created_at ASC", (parent_id,)
+            ).fetchall()
+            return [
+                Task(
+                    task_id=row["task_id"],
+                    prompt=row["prompt"],
+                    task_type=row["task_type"],
+                    status=TaskStatus(row["status"]),
+                    progress=row["progress"],
+                    step=row["step"],
+                    result=row["result"],
+                    error=row["error"],
+                    parent_id=row["parent_id"],
+                    created_at=row["created_at"],
+                    updated_at=row["updated_at"],
+                )
+                for row in rows
+            ]
+
+
+class PostgreSQLTaskStore(BaseTaskStore):
+    """PostgreSQL implementation of the TaskStore interface using connection pooling."""
+    
+    _pool = None
+    
+    def __init__(self, connection_url: Optional[str] = None):
+        self.connection_url = connection_url
+        if not self.connection_url:
+            from agyqueue.config import settings
+            self.connection_url = settings.database_url
+            
+        self._init_pool()
+        self._init_db()
+
+    def _init_pool(self):
+        if PostgreSQLTaskStore._pool is None:
+            import psycopg2.pool
+            PostgreSQLTaskStore._pool = psycopg2.pool.ThreadedConnectionPool(
+                minconn=1,
+                maxconn=10,
+                dsn=self.connection_url
+            )
+            logger.info("PostgreSQL thread pool connection manager initialized.")
+
+    @contextmanager
+    def _get_conn(self):
+        conn = PostgreSQLTaskStore._pool.getconn()
+        conn.autocommit = True
+        try:
+            yield conn
+        finally:
+            PostgreSQLTaskStore._pool.putconn(conn)
+
+    def _init_db(self):
+        with self._get_conn() as conn:
+            with conn.cursor() as cur:
+                cur.execute("""
+                    CREATE TABLE IF NOT EXISTS tasks (
+                        task_id VARCHAR(50) PRIMARY KEY,
+                        prompt TEXT NOT NULL,
+                        task_type VARCHAR(50) NOT NULL,
+                        status VARCHAR(20) NOT NULL,
+                        progress INTEGER NOT NULL,
+                        step TEXT NOT NULL,
+                        result TEXT,
+                        error TEXT,
+                        parent_id VARCHAR(50),
+                        created_at VARCHAR(30) NOT NULL,
+                        updated_at VARCHAR(30) NOT NULL
+                    )
+                """)
+
+    def save_task(self, task: Task) -> None:
+        with self._get_conn() as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    """
+                    INSERT INTO tasks 
+                    (task_id, prompt, task_type, status, progress, step, result, error, parent_id, created_at, updated_at)
+                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
+                    ON CONFLICT (task_id) DO UPDATE SET
+                        prompt = EXCLUDED.prompt,
+                        task_type = EXCLUDED.task_type,
+                        status = EXCLUDED.status,
+                        progress = EXCLUDED.progress,
+                        step = EXCLUDED.step,
+                        result = EXCLUDED.result,
+                        error = EXCLUDED.error,
+                        parent_id = EXCLUDED.parent_id,
+                        created_at = EXCLUDED.created_at,
+                        updated_at = EXCLUDED.updated_at
+                    """,
+                    (
+                        task.task_id,
+                        task.prompt,
+                        task.task_type,
+                        task.status.value,
+                        task.progress,
+                        task.step,
+                        task.result,
+                        task.error,
+                        task.parent_id,
+                        task.created_at,
+                        task.updated_at,
+                    ),
+                )
+
+    def get_task(self, task_id: str) -> Optional[Task]:
+        from psycopg2.extras import RealDictCursor
+        with self._get_conn() as conn:
+            with conn.cursor(cursor_factory=RealDictCursor) as cur:
+                cur.execute("SELECT * FROM tasks WHERE task_id = %s", (task_id,))
+                row = cur.fetchone()
+                if not row:
+                    return None
+                return Task(
+                    task_id=row["task_id"],
+                    prompt=row["prompt"],
+                    task_type=row["task_type"],
+                    status=TaskStatus(row["status"]),
+                    progress=row["progress"],
+                    step=row["step"],
+                    result=row["result"],
+                    error=row["error"],
+                    parent_id=row["parent_id"],
+                    created_at=row["created_at"],
+                    updated_at=row["updated_at"],
+                )
+
+    def update_task(
+        self,
+        task_id: str,
+        status: TaskStatus,
+        progress: int,
+        step: str,
+        result: Optional[str] = None,
+        error: Optional[str] = None,
+    ) -> Optional[Task]:
+        now = datetime.now(timezone.utc).replace(tzinfo=None).isoformat()
+        with self._get_conn() as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    """
+                    UPDATE tasks
+                    SET status = %s, progress = %s, step = %s, result = COALESCE(%s, result), error = COALESCE(%s, error), updated_at = %s
+                    WHERE task_id = %s
+                    """,
+                    (status.value, progress, step, result, error, now, task_id),
+                )
+        return self.get_task(task_id)
+
+    def touch_task(self, task_id: str) -> None:
+        now = datetime.now(timezone.utc).replace(tzinfo=None).isoformat()
+        with self._get_conn() as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    "UPDATE tasks SET updated_at = %s WHERE task_id = %s",
+                    (now, task_id),
+                )
+
+    def list_tasks(self) -> List[Task]:
+        from psycopg2.extras import RealDictCursor
+        with self._get_conn() as conn:
+            with conn.cursor(cursor_factory=RealDictCursor) as cur:
+                cur.execute("SELECT * FROM tasks ORDER BY created_at DESC")
+                rows = cur.fetchall()
+                return [
+                    Task(
+                        task_id=row["task_id"],
+                        prompt=row["prompt"],
+                        task_type=row["task_type"],
+                        status=TaskStatus(row["status"]),
+                        progress=row["progress"],
+                        step=row["step"],
+                        result=row["result"],
+                        error=row["error"],
+                        parent_id=row["parent_id"],
+                        created_at=row["created_at"],
+                        updated_at=row["updated_at"],
+                    )
+                    for row in rows
+                ]
+
+    def get_subtasks(self, parent_id: str) -> List[Task]:
+        from psycopg2.extras import RealDictCursor
+        with self._get_conn() as conn:
+            with conn.cursor(cursor_factory=RealDictCursor) as cur:
+                cur.execute(
+                    "SELECT * FROM tasks WHERE parent_id = %s ORDER BY created_at ASC", (parent_id,)
+                )
+                rows = cur.fetchall()
+                return [
+                    Task(
+                        task_id=row["task_id"],
+                        prompt=row["prompt"],
+                        task_type=row["task_type"],
+                        status=TaskStatus(row["status"]),
+                        progress=row["progress"],
+                        step=row["step"],
+                        result=row["result"],
+                        error=row["error"],
+                        parent_id=row["parent_id"],
+                        created_at=row["created_at"],
+                        updated_at=row["updated_at"],
+                    )
+                    for row in rows
+                ]
+
+
+def TaskStore(db_path: Optional[str] = None) -> BaseTaskStore:
+    """Factory function returning the configured TaskStore implementation.
+    
+    Can be configured via environment variables (e.g., AGYQUEUE_STORE_TYPE).
+    """
+    from agyqueue.config import settings
+    
+    store_type = settings.store_type
+    
+    # If postgres is set or database_url starts with postgres, use PostgreSQLTaskStore
+    if store_type == "postgres" or (settings.database_url and settings.database_url.startswith("postgres")):
+        try:
+            return PostgreSQLTaskStore()
+        except ImportError:
+            logger.warning("psycopg2 not installed. Falling back to SQLiteTaskStore.")
+            return SQLiteTaskStore(db_path)
+        except Exception as e:
+            logger.error(f"Failed to connect to PostgreSQL: {e}. Falling back to SQLiteTaskStore.")
+            return SQLiteTaskStore(db_path)
+            
+    return SQLiteTaskStore(db_path)

agyqueue/task_queue.py

@@ -0,0 +1,111 @@
+import os
+import time
+import logging
+from abc import ABC, abstractmethod
+from typing import Optional
+import redis
+from agyqueue.storage import TaskStore
+from agyqueue.models import TaskStatus
+
+logger = logging.getLogger(__name__)
+
+class BaseTaskQueue(ABC):
+    """Abstract base class defining the interface for task queues."""
+    
+    @abstractmethod
+    def enqueue(self, task_id: str) -> None:
+        """Pushes a task ID onto the queue."""
+        pass
+
+    @abstractmethod
+    def dequeue(self, timeout: int = 1) -> Optional[str]:
+        """Pops a task ID from the queue. Blocks up to `timeout` seconds if empty."""
+        pass
+
+
+class RedisTaskQueue(BaseTaskQueue):
+    """Redis-backed implementation of the TaskQueue interface."""
+    
+    def __init__(self, redis_url: str):
+        self.redis_url = redis_url
+        self.redis_client = redis.Redis.from_url(
+            self.redis_url, 
+            socket_connect_timeout=2.0,
+            decode_responses=True
+        )
+        # Test connection
+        self.redis_client.ping()
+        logger.info(f"Connected to Redis queue at {self.redis_url}")
+
+    def enqueue(self, task_id: str) -> None:
+        self.redis_client.rpush("agyqueue:task_ids", task_id)
+        logger.info(f"Enqueued task {task_id} to Redis.")
+
+    def dequeue(self, timeout: int = 1) -> Optional[str]:
+        # blpop returns (queue_name, value)
+        res = self.redis_client.blpop("agyqueue:task_ids", timeout=timeout)
+        if res:
+            return res[1]
+        return None
+
+
+class SQLiteTaskQueue(BaseTaskQueue):
+    """SQLite-backed polling implementation of the TaskQueue interface."""
+    
+    def __init__(self, db_path: Optional[str] = None):
+        self.db_path = db_path or os.environ.get("AGYQUEUE_DB_PATH", "agyqueue.db")
+        self.store = TaskStore(self.db_path)
+
+    def enqueue(self, task_id: str) -> None:
+        # SQLite queue relies on the task status being set to QUEUED in the SQLite store
+        logger.info(f"Enqueued task {task_id} to SQLite store (waiting for polling worker).")
+
+    def dequeue(self, timeout: int = 1) -> Optional[str]:
+        try:
+            # We access the internal connection from the sqlite store to perform the dequeue transaction
+            # Since the store uses WAL mode, we can claim the task concurrently.
+            # Wait, since self.store is a SQLiteTaskStore instance, we can call _get_conn()
+            # If the store is abstract, we verify if it has _get_conn
+            if hasattr(self.store, "_get_conn"):
+                conn = self.store._get_conn()
+                with conn:
+                    row = conn.execute(
+                        "SELECT task_id FROM tasks WHERE status = ? ORDER BY created_at ASC LIMIT 1",
+                        (TaskStatus.QUEUED.value,)
+                    ).fetchone()
+                    
+                    if row:
+                        task_id = row["task_id"]
+                        # Claim the task
+                        conn.execute(
+                            "UPDATE tasks SET status = ?, step = ?, updated_at = ? WHERE task_id = ?",
+                            (
+                                TaskStatus.RUNNING.value,
+                                "Claimed by worker",
+                                time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+                                task_id
+                            )
+                        )
+                        return task_id
+        except Exception as e:
+            logger.error(f"SQLite dequeue failed: {e}")
+        
+        time.sleep(timeout)
+        return None
+
+
+def TaskQueue(redis_url: Optional[str] = None, db_path: Optional[str] = None) -> BaseTaskQueue:
+    """Factory function returning the configured TaskQueue implementation.
+    
+    Tries to connect to Redis if a URL is provided or set in environment,
+    falling back to SQLite if Redis is unavailable.
+    """
+    url = redis_url or os.environ.get("REDIS_URL")
+    if url:
+        try:
+            return RedisTaskQueue(url)
+        except Exception as e:
+            logger.warning(f"Could not connect to Redis at {url}: {e}. Falling back to SQLite queue.")
+    
+    logger.info("Using SQLite database as task queue.")
+    return SQLiteTaskQueue(db_path)