| import json |
| import os |
| import pathlib |
| import time |
| import logging |
| from typing import Any |
|
|
| logger = logging.getLogger("qkernel.queue") |
|
|
| DEFAULT_DB = pathlib.Path(__file__).resolve().parent.parent / "data" / "queue.db" |
|
|
|
|
| class PersistentQueue: |
| def __init__(self, db_path: str | pathlib.Path | None = None) -> None: |
| self.db_path = pathlib.Path(db_path or os.getenv("QUEUE_DB_PATH", DEFAULT_DB)) |
| self.db_path.parent.mkdir(parents=True, exist_ok=True) |
| self._conn = self._connect() |
| self._init_db() |
|
|
| def _connect(self): |
| import sqlite3 |
|
|
| conn = sqlite3.connect(str(self.db_path), check_same_thread=False) |
| conn.row_factory = sqlite3.Row |
| return conn |
|
|
| def _init_db(self) -> None: |
| self._conn.executescript( |
| """ |
| CREATE TABLE IF NOT EXISTS queue ( |
| id TEXT PRIMARY KEY, |
| type TEXT NOT NULL, |
| payload TEXT NOT NULL, |
| status TEXT NOT NULL, |
| created_at INTEGER NOT NULL, |
| updated_at INTEGER NOT NULL |
| ); |
| CREATE INDEX IF NOT EXISTS idx_queue_status ON queue(status, created_at); |
| """ |
| ) |
| self._conn.commit() |
|
|
| def enqueue(self, type_: str, payload: dict) -> str: |
| import uuid |
|
|
| row_id = f"q_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}" |
| now = int(time.time() * 1000) |
| self._conn.execute( |
| "INSERT INTO queue (id, type, payload, status, created_at, updated_at) " |
| "VALUES (?, ?, ?, 'pending', ?, ?)", |
| (row_id, type_, json.dumps(payload, ensure_ascii=False), now, now), |
| ) |
| self._conn.commit() |
| return row_id |
|
|
| def dequeue(self) -> dict[str, Any] | None: |
| tx = self._conn.execute( |
| "SELECT * FROM queue WHERE status = 'pending' " |
| "ORDER BY created_at LIMIT 1" |
| ).fetchone() |
| if tx is None: |
| return None |
| row = { |
| "id": tx["id"], |
| "type": tx["type"], |
| "payload": json.loads(tx["payload"]), |
| "status": tx["status"], |
| "created_at": tx["created_at"], |
| "updated_at": tx["updated_at"], |
| } |
| self._conn.execute( |
| "UPDATE queue SET status = 'processing', updated_at = ? WHERE id = ?", |
| (int(time.time() * 1000), row["id"]), |
| ) |
| self._conn.commit() |
| return row |
|
|
| def complete(self, task_id: str, result: dict | None = None) -> None: |
| self._conn.execute( |
| "UPDATE queue SET status = 'completed', updated_at = ?, payload = ? " |
| "WHERE id = ?", |
| ( |
| int(time.time() * 1000), |
| json.dumps({"result": result}, ensure_ascii=False), |
| task_id, |
| ), |
| ) |
| self._conn.commit() |
|
|
| def fail(self, task_id: str, error: str) -> None: |
| self._conn.execute( |
| "UPDATE queue SET status = 'failed', updated_at = ? , payload = ? " |
| "WHERE id = ?", |
| ( |
| int(time.time() * 1000), |
| json.dumps({"error": error}, ensure_ascii=False), |
| task_id, |
| ), |
| ) |
| self._conn.commit() |
|
|
| def requeue(self, task_id: str) -> None: |
| self._conn.execute( |
| "UPDATE queue SET status = 'pending', updated_at = ? WHERE id = ?", |
| (int(time.time() * 1000), task_id), |
| ) |
| self._conn.commit() |
|
|
| def get(self, task_id: str) -> dict[str, Any] | None: |
| tx = self._conn.execute( |
| "SELECT * FROM queue WHERE id = ?", (task_id,) |
| ).fetchone() |
| if tx is None: |
| return None |
| return { |
| "id": tx["id"], |
| "type": tx["type"], |
| "payload": json.loads(tx["payload"]), |
| "status": tx["status"], |
| "created_at": tx["created_at"], |
| "updated_at": tx["updated_at"], |
| } |
|
|
| def stats(self) -> dict: |
| rows = self._conn.execute( |
| "SELECT status, COUNT(*) AS count FROM queue GROUP BY status" |
| ).fetchall() |
| return {r["status"]: r["count"] for r in rows} or {"pending": 0} |
|
|
| def clear(self, status: str | None = None) -> int: |
| if status: |
| cur = self._conn.execute( |
| "DELETE FROM queue WHERE status = ?", (status,) |
| ) |
| else: |
| cur = self._conn.execute("DELETE FROM queue") |
| self._conn.commit() |
| return cur.rowcount |
|
|
| def close(self) -> None: |
| self._conn.close() |
