52 lines
2.0 KiB
Python
52 lines
2.0 KiB
Python
import sqlite3
|
|
import json
|
|
from typing import Any, Dict, List
|
|
from core.config import Env
|
|
|
|
class SqliteQueue:
|
|
"""
|
|
"""
|
|
|
|
def __init__(self, db_path: str = Env.SQLITE_PATH):
|
|
self.db_path = db_path
|
|
self._init_schema()
|
|
|
|
def _conn(self):
|
|
conn = sqlite3.connect(self.db_path)
|
|
conn.row_factory = sqlite3.Row
|
|
conn.execute("PRAGMA journal_mode=WAL;")
|
|
conn.execute("PRAGMA synchronous=NORMAL;")
|
|
return conn
|
|
|
|
def _init_schema(self):
|
|
with self._conn() as con:
|
|
con.executescript("""
|
|
CREATE TABLE IF NOT EXISTS tasks(
|
|
task_id TEXT PRIMARY KEY,
|
|
queue TEXT NOT NULL,
|
|
type TEXT NOT NULL,
|
|
payload_json TEXT NOT NULL,
|
|
created_at INTEGER NOT NULL,
|
|
status TEXT DEFAULT 'pending',
|
|
attempts INTEGER DEFAULT 0,
|
|
last_error TEXT
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_tasks_queue_status ON tasks(queue, status);
|
|
""")
|
|
|
|
def add_task(self, task: Dict[str, Any]) -> None:
|
|
with self._conn() as con:
|
|
con.execute("""
|
|
INSERT OR REPLACE INTO tasks(task_id, queue, type, payload_json, created_at, status, attempts)
|
|
VALUES(?,?,?,?,?,'pending',?)
|
|
""", (task["task_id"], task["queue"], task["type"], json.dumps(task["payload"]), task["created_at"], task.get("_attempts", 0)))
|
|
|
|
def update_task(self, task_id: str, status: str, error: str | None = None, attempts: int | None = None) -> None:
|
|
with self._conn() as con:
|
|
con.execute("""UPDATE tasks SET status=?, last_error=?, attempts=COALESCE(?, attempts) WHERE task_id=? """, (status, error, attempts, task_id))
|
|
|
|
def pending_for_queue(self, queue: str) -> List[Dict[str, Any]]:
|
|
with self._conn() as con:
|
|
rows = con.execute("SELECT * FROM tasks WHERE queue=? AND status='pending'", (queue,)).fetchall()
|
|
return [dict(r) for r in rows]
|