File size: 4,585 Bytes
945aae0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | 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() |