baqueue 0.1.0__py3-none-any.whl

baqueue/drivers/postgres_driver.py

@@ -0,0 +1,656 @@
+"""PostgreSQL driver for BaQueue - reliable transactional queue backend."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from baqueue.drivers.base import BaseDriver
+from baqueue.serializer import JobPayload, _now_ts
+
+
+class PostgresDriver(BaseDriver):
+    """PostgreSQL-backed driver using a single jobs table with advisory locks.
+
+    Creates two tables:
+        {prefix}_jobs    - all job data
+        {prefix}_batches - batch metadata
+        {prefix}_metrics - metric entries
+    """
+
+    def __init__(self, url: str = "", prefix: str = "baqueue", **kwargs: Any):
+        self._url = url
+        self._prefix = prefix
+        self._kwargs = kwargs
+        self._pool: Any = None
+
+    # Postgres sqlstate codes for storage exhaustion. asyncpg surfaces these on
+    # the exception's `sqlstate` attribute.
+    _STORAGE_SQLSTATES = frozenset({"53100", "53200", "53300", "54000"})
+
+    def is_storage_full_error(self, exc: BaseException) -> bool:
+        sqlstate = getattr(exc, "sqlstate", None) or getattr(exc, "pgcode", None)
+        if sqlstate in self._STORAGE_SQLSTATES:
+            return True
+        msg = str(exc).lower()
+        return (
+            "disk full" in msg
+            or "no space left" in msg
+            or "out of memory" in msg
+            or "could not extend file" in msg
+        )
+
+    @property
+    def _jobs_table(self) -> str:
+        return f"{self._prefix}_jobs"
+
+    @property
+    def _batches_table(self) -> str:
+        return f"{self._prefix}_batches"
+
+    @property
+    def _metrics_table(self) -> str:
+        return f"{self._prefix}_metrics"
+
+    @property
+    def _supervisors_table(self) -> str:
+        return f"{self._prefix}_supervisors"
+
+    async def connect(self) -> None:
+        try:
+            import asyncpg
+        except ImportError:
+            raise ImportError(
+                "PostgreSQL driver requires the 'asyncpg' package.\n"
+                "Install it with: pip install baqueue[postgres]"
+            ) from None
+        self._pool = await asyncpg.create_pool(self._url, **self._kwargs)
+        await self._ensure_tables()
+
+    async def disconnect(self) -> None:
+        if self._pool:
+            await self._pool.close()
+            self._pool = None
+
+    async def _ensure_tables(self) -> None:
+        async with self._pool.acquire() as conn:
+            await conn.execute(f"""
+                CREATE TABLE IF NOT EXISTS {self._jobs_table} (
+                    id TEXT PRIMARY KEY,
+                    job_class TEXT NOT NULL,
+                    data JSONB NOT NULL DEFAULT '{{}}',
+                    queue TEXT NOT NULL DEFAULT 'default',
+                    status TEXT NOT NULL DEFAULT 'pending',
+                    attempts INTEGER NOT NULL DEFAULT 0,
+                    max_attempts INTEGER NOT NULL DEFAULT 3,
+                    backoff TEXT NOT NULL DEFAULT 'exponential',
+                    timeout INTEGER NOT NULL DEFAULT 60,
+                    tags TEXT[] NOT NULL DEFAULT '{{}}',
+                    batch_id TEXT,
+                    delay_until DOUBLE PRECISION,
+                    error TEXT,
+                    created_at DOUBLE PRECISION NOT NULL,
+                    updated_at DOUBLE PRECISION NOT NULL,
+                    started_at DOUBLE PRECISION,
+                    completed_at DOUBLE PRECISION,
+                    failed_at DOUBLE PRECISION
+                )
+            """)
+            await conn.execute(f"""
+                CREATE INDEX IF NOT EXISTS idx_{self._prefix}_queue_status_created
+                ON {self._jobs_table} (queue, status, created_at DESC)
+            """)
+            await conn.execute(f"""
+                CREATE INDEX IF NOT EXISTS idx_{self._prefix}_created_at
+                ON {self._jobs_table} (created_at DESC)
+            """)
+            await conn.execute(f"""
+                CREATE INDEX IF NOT EXISTS idx_{self._prefix}_status_created
+                ON {self._jobs_table} (status, created_at DESC)
+            """)
+            await conn.execute(f"""
+                CREATE INDEX IF NOT EXISTS idx_{self._prefix}_delay
+                ON {self._jobs_table} (delay_until) WHERE delay_until IS NOT NULL
+            """)
+            await conn.execute(f"""
+                CREATE INDEX IF NOT EXISTS idx_{self._prefix}_batch
+                ON {self._jobs_table} (batch_id) WHERE batch_id IS NOT NULL
+            """)
+            await conn.execute(f"DROP INDEX IF EXISTS idx_{self._prefix}_queue_status")
+            await conn.execute(f"""
+                CREATE TABLE IF NOT EXISTS {self._batches_table} (
+                    id TEXT PRIMARY KEY,
+                    data JSONB NOT NULL DEFAULT '{{}}'
+                )
+            """)
+            await conn.execute(f"""
+                CREATE TABLE IF NOT EXISTS {self._metrics_table} (
+                    id SERIAL PRIMARY KEY,
+                    queue TEXT NOT NULL,
+                    metric TEXT NOT NULL,
+                    value DOUBLE PRECISION NOT NULL,
+                    recorded_at DOUBLE PRECISION NOT NULL
+                )
+            """)
+            await conn.execute(f"""
+                CREATE INDEX IF NOT EXISTS idx_{self._prefix}_metrics_queue_metric
+                ON {self._metrics_table} (queue, metric)
+            """)
+            await conn.execute(f"""
+                CREATE INDEX IF NOT EXISTS idx_{self._prefix}_metrics_recorded_at
+                ON {self._metrics_table} (recorded_at)
+            """)
+            await conn.execute(f"""
+                CREATE TABLE IF NOT EXISTS {self._supervisors_table} (
+                    name TEXT PRIMARY KEY,
+                    data JSONB NOT NULL DEFAULT '{{}}',
+                    heartbeat_at DOUBLE PRECISION NOT NULL
+                )
+            """)
+            await conn.execute(f"""
+                CREATE INDEX IF NOT EXISTS idx_{self._prefix}_supervisors_heartbeat
+                ON {self._supervisors_table} (heartbeat_at)
+            """)
+
+    def _row_to_payload(self, row: Any) -> JobPayload:
+        backoff: str | list[int] = row["backoff"]
+        try:
+            parsed = json.loads(backoff)
+            if isinstance(parsed, list):
+                backoff = parsed
+        except (json.JSONDecodeError, TypeError):
+            pass
+
+        return JobPayload(
+            id=row["id"],
+            job_class=row["job_class"],
+            data=json.loads(row["data"]) if isinstance(row["data"], str) else row["data"],
+            queue=row["queue"],
+            status=row["status"],
+            attempts=row["attempts"],
+            max_attempts=row["max_attempts"],
+            backoff=backoff,
+            timeout=row["timeout"],
+            tags=list(row["tags"]) if row["tags"] else [],
+            batch_id=row["batch_id"],
+            delay_until=row["delay_until"],
+            error=row["error"],
+            created_at=row["created_at"],
+            updated_at=row["updated_at"],
+            started_at=row["started_at"],
+            completed_at=row["completed_at"],
+            failed_at=row["failed_at"],
+        )
+
+    # ── Push / Pop ──────────────────────────────────────────────
+
+    async def push(self, payload: JobPayload) -> str:
+        payload.status = "pending"
+        payload.updated_at = _now_ts()
+        backoff_str = json.dumps(payload.backoff) if isinstance(payload.backoff, list) else payload.backoff
+
+        async def _do():
+            async with self._pool.acquire() as conn:
+                await conn.execute(
+                    f"""INSERT INTO {self._jobs_table}
+                        (id, job_class, data, queue, status, attempts, max_attempts,
+                         backoff, timeout, tags, batch_id, delay_until, created_at, updated_at)
+                        VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14)""",
+                    payload.id, payload.job_class, json.dumps(payload.data),
+                    payload.queue, payload.status, payload.attempts, payload.max_attempts,
+                    backoff_str, payload.timeout, payload.tags,
+                    payload.batch_id, payload.delay_until,
+                    payload.created_at, payload.updated_at,
+                )
+        await self._with_disk_full_recovery(_do)
+        return payload.id
+
+    async def push_many(self, payloads: list[JobPayload]) -> list[str]:
+        ids: list[str] = []
+        now = _now_ts()
+
+        async def _do():
+            ids.clear()
+            async with self._pool.acquire() as conn:
+                async with conn.transaction():
+                    for p in payloads:
+                        p.status = "pending"
+                        p.updated_at = now
+                        backoff_str = json.dumps(p.backoff) if isinstance(p.backoff, list) else p.backoff
+                        await conn.execute(
+                            f"""INSERT INTO {self._jobs_table}
+                                (id, job_class, data, queue, status, attempts, max_attempts,
+                                 backoff, timeout, tags, batch_id, delay_until, created_at, updated_at)
+                                VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14)""",
+                            p.id, p.job_class, json.dumps(p.data),
+                            p.queue, p.status, p.attempts, p.max_attempts,
+                            backoff_str, p.timeout, p.tags,
+                            p.batch_id, p.delay_until,
+                            p.created_at, p.updated_at,
+                        )
+                        ids.append(p.id)
+        await self._with_disk_full_recovery(_do)
+        return ids
+
+    async def pop(self, queue: str) -> JobPayload | None:
+        now = _now_ts()
+
+        async def _do():
+            async with self._pool.acquire() as conn:
+                return await conn.fetchrow(
+                    f"""UPDATE {self._jobs_table}
+                        SET status='processing', started_at=$1, updated_at=$1, attempts=attempts+1
+                        WHERE id = (
+                            SELECT id FROM {self._jobs_table}
+                            WHERE queue=$2 AND status='pending'
+                              AND (delay_until IS NULL OR delay_until <= $1)
+                            ORDER BY created_at ASC
+                            FOR UPDATE SKIP LOCKED
+                            LIMIT 1
+                        )
+                        RETURNING *""",
+                    now, queue,
+                )
+
+        row = await self._with_disk_full_recovery(_do)
+        if row:
+            return self._row_to_payload(row)
+        return None
+
+    async def pop_delayed(self) -> list[JobPayload]:
+        now = _now_ts()
+
+        async def _do():
+            async with self._pool.acquire() as conn:
+                return await conn.fetch(
+                    f"""UPDATE {self._jobs_table}
+                        SET delay_until=NULL, updated_at=$1
+                        WHERE status='pending' AND delay_until IS NOT NULL AND delay_until <= $1
+                        RETURNING *""",
+                    now,
+                )
+
+        rows = await self._with_disk_full_recovery(_do)
+        return [self._row_to_payload(r) for r in rows]
+
+    # ── Job lifecycle ───────────────────────────────────────────
+
+    async def complete(self, payload: JobPayload) -> None:
+        now = _now_ts()
+
+        async def _do():
+            async with self._pool.acquire() as conn:
+                await conn.execute(
+                    f"UPDATE {self._jobs_table} SET status='completed', completed_at=$1, updated_at=$1 WHERE id=$2",
+                    now, payload.id,
+                )
+
+        await self._with_disk_full_recovery(_do)
+        payload.status = "completed"
+        payload.completed_at = now
+
+    async def fail(self, payload: JobPayload, error: str) -> None:
+        now = _now_ts()
+
+        async def _do():
+            async with self._pool.acquire() as conn:
+                await conn.execute(
+                    f"UPDATE {self._jobs_table} SET status='failed', failed_at=$1, updated_at=$1, error=$2 WHERE id=$3",
+                    now, error, payload.id,
+                )
+
+        await self._with_disk_full_recovery(_do)
+        payload.status = "failed"
+        payload.failed_at = now
+        payload.error = error
+
+    async def release(self, payload: JobPayload, delay: float = 0) -> None:
+        now = _now_ts()
+        delay_until = now + delay if delay > 0 else None
+
+        async def _do():
+            async with self._pool.acquire() as conn:
+                await conn.execute(
+                    f"UPDATE {self._jobs_table} SET status='pending', updated_at=$1, delay_until=$2 WHERE id=$3",
+                    now, delay_until, payload.id,
+                )
+
+        await self._with_disk_full_recovery(_do)
+
+    async def delete(self, job_id: str) -> None:
+        async def _do():
+            async with self._pool.acquire() as conn:
+                await conn.execute(f"DELETE FROM {self._jobs_table} WHERE id=$1", job_id)
+
+        await self._with_disk_full_recovery(_do)
+
+    # ── Query ───────────────────────────────────────────────────
+
+    async def get_job(self, job_id: str) -> JobPayload | None:
+        async with self._pool.acquire() as conn:
+            row = await conn.fetchrow(f"SELECT * FROM {self._jobs_table} WHERE id=$1", job_id)
+        return self._row_to_payload(row) if row else None
+
+    def _build_where(
+        self,
+        queue: str | None = None,
+        status: str | None = None,
+        tag: str | None = None,
+        batch_id: str | None = None,
+        created_from: float | None = None,
+        created_to: float | None = None,
+        start_idx: int = 1,
+    ) -> tuple[str, list[Any], int]:
+        conditions: list[str] = []
+        params: list[Any] = []
+        idx = start_idx
+
+        if queue:
+            conditions.append(f"queue=${idx}")
+            params.append(queue)
+            idx += 1
+        if status:
+            conditions.append(f"status=${idx}")
+            params.append(status)
+            idx += 1
+        if tag:
+            conditions.append(f"${idx} = ANY(tags)")
+            params.append(tag)
+            idx += 1
+        if batch_id:
+            conditions.append(f"batch_id=${idx}")
+            params.append(batch_id)
+            idx += 1
+        if created_from is not None:
+            conditions.append(f"created_at >= ${idx}")
+            params.append(created_from)
+            idx += 1
+        if created_to is not None:
+            conditions.append(f"created_at <= ${idx}")
+            params.append(created_to)
+            idx += 1
+
+        where = " AND ".join(conditions) if conditions else "TRUE"
+        return where, params, idx
+
+    async def get_jobs(
+        self,
+        queue: str | None = None,
+        status: str | None = None,
+        tag: str | None = None,
+        batch_id: str | None = None,
+        offset: int = 0,
+        limit: int = 50,
+        created_from: float | None = None,
+        created_to: float | None = None,
+    ) -> list[JobPayload]:
+        where, params, idx = self._build_where(
+            queue, status, tag, batch_id, created_from, created_to,
+        )
+        params.extend([limit, offset])
+
+        async with self._pool.acquire() as conn:
+            rows = await conn.fetch(
+                f"SELECT * FROM {self._jobs_table} WHERE {where} ORDER BY created_at DESC LIMIT ${idx} OFFSET ${idx+1}",
+                *params,
+            )
+        return [self._row_to_payload(r) for r in rows]
+
+    async def count_jobs(
+        self,
+        queue: str | None = None,
+        status: str | None = None,
+        created_from: float | None = None,
+        created_to: float | None = None,
+    ) -> int:
+        where, params, _ = self._build_where(
+            queue=queue, status=status,
+            created_from=created_from, created_to=created_to,
+        )
+        async with self._pool.acquire() as conn:
+            row = await conn.fetchrow(
+                f"SELECT COUNT(*) AS cnt FROM {self._jobs_table} WHERE {where}",
+                *params,
+            )
+        return int(row["cnt"]) if row else 0
+
+    async def size(self, queue: str) -> int:
+        async with self._pool.acquire() as conn:
+            row = await conn.fetchrow(
+                f"SELECT COUNT(*) as cnt FROM {self._jobs_table} WHERE queue=$1 AND status='pending'",
+                queue,
+            )
+        return row["cnt"] if row else 0
+
+    async def queues(self) -> list[str]:
+        async with self._pool.acquire() as conn:
+            rows = await conn.fetch(f"SELECT DISTINCT queue FROM {self._jobs_table} ORDER BY queue")
+        return [r["queue"] for r in rows]
+
+    # ── Metrics ─────────────────────────────────────────────────
+
+    async def record_metric(self, queue: str, metric: str, value: float) -> None:
+        async def _do():
+            async with self._pool.acquire() as conn:
+                await conn.execute(
+                    f"INSERT INTO {self._metrics_table} (queue, metric, value, recorded_at) VALUES ($1,$2,$3,$4)",
+                    queue, metric, value, _now_ts(),
+                )
+        await self._with_disk_full_recovery(_do)
+
+    async def get_metrics(self, queue: str | None = None) -> dict[str, Any]:
+        """Live status counts from the jobs table — never from the metrics event log."""
+        async with self._pool.acquire() as conn:
+            if queue:
+                rows = await conn.fetch(
+                    f"SELECT queue, status, COUNT(*) AS cnt FROM {self._jobs_table} "
+                    f"WHERE queue=$1 GROUP BY queue, status",
+                    queue,
+                )
+                queue_names = [queue]
+            else:
+                rows = await conn.fetch(
+                    f"SELECT queue, status, COUNT(*) AS cnt FROM {self._jobs_table} GROUP BY queue, status"
+                )
+                qrows = await conn.fetch(
+                    f"SELECT DISTINCT queue FROM {self._jobs_table} ORDER BY queue"
+                )
+                queue_names = [r["queue"] for r in qrows]
+
+        result: dict[str, Any] = {
+            q: {"pending": 0, "processing": 0, "completed": 0, "failed": 0}
+            for q in queue_names
+        }
+        for r in rows:
+            q = r["queue"]
+            if q not in result:
+                result[q] = {"pending": 0, "processing": 0, "completed": 0, "failed": 0}
+            if r["status"] in result[q]:
+                result[q][r["status"]] = int(r["cnt"])
+        return result
+
+    async def report_supervisor(self, stats: dict[str, Any]) -> None:
+        name = str(stats.get("name", "")).strip()
+        if not name:
+            return
+        now = _now_ts()
+        payload = json.dumps(stats)
+
+        async def _do():
+            async with self._pool.acquire() as conn:
+                await conn.execute(
+                    f"""INSERT INTO {self._supervisors_table} (name, data, heartbeat_at)
+                        VALUES ($1, $2::jsonb, $3)
+                        ON CONFLICT (name) DO UPDATE SET
+                            data=EXCLUDED.data,
+                            heartbeat_at=EXCLUDED.heartbeat_at""",
+                    name, payload, now,
+                )
+
+        await self._with_disk_full_recovery(_do)
+
+    async def get_supervisor_stats(self, stale_after: float = 10.0) -> list[dict[str, Any]]:
+        cutoff = _now_ts() - stale_after
+        async with self._pool.acquire() as conn:
+            rows = await conn.fetch(
+                f"""SELECT data FROM {self._supervisors_table}
+                    WHERE heartbeat_at >= $1
+                    ORDER BY name""",
+                cutoff,
+            )
+        out: list[dict[str, Any]] = []
+        for row in rows:
+            raw = row["data"]
+            data = json.loads(raw) if isinstance(raw, str) else dict(raw)
+            if not isinstance(data, dict):
+                continue
+            if not data.get("running", False):
+                continue
+            out.append(data)
+        return out
+
+    # ── Batch helpers ───────────────────────────────────────────
+
+    async def store_batch(self, batch_id: str, data: dict[str, Any]) -> None:
+        async def _do():
+            async with self._pool.acquire() as conn:
+                await conn.execute(
+                    f"INSERT INTO {self._batches_table} (id, data) VALUES ($1, $2)",
+                    batch_id, json.dumps(data),
+                )
+
+        await self._with_disk_full_recovery(_do)
+
+    async def get_batch(self, batch_id: str) -> dict[str, Any] | None:
+        async with self._pool.acquire() as conn:
+            row = await conn.fetchrow(f"SELECT data FROM {self._batches_table} WHERE id=$1", batch_id)
+        if row:
+            return json.loads(row["data"]) if isinstance(row["data"], str) else row["data"]
+        return None
+
+    async def update_batch(self, batch_id: str, data: dict[str, Any]) -> None:
+        async def _do():
+            async with self._pool.acquire() as conn:
+                await conn.execute(
+                    f"UPDATE {self._batches_table} SET data=$1 WHERE id=$2",
+                    json.dumps(data), batch_id,
+                )
+
+        await self._with_disk_full_recovery(_do)
+
+    async def increment_batch_counter(
+        self, batch_id: str, field: str, delta: int = 1,
+    ) -> dict[str, Any] | None:
+        # jsonb_set with COALESCE so missing fields start at 0. Returns the
+        # post-update row in a single statement -> atomic.
+        async def _do():
+            async with self._pool.acquire() as conn:
+                return await conn.fetchrow(
+                    f"""UPDATE {self._batches_table}
+                        SET data = jsonb_set(
+                            data,
+                            ARRAY[$1],
+                            to_jsonb(COALESCE((data->>$1)::int, 0) + $2)
+                        )
+                        WHERE id = $3
+                        RETURNING data""",
+                    field, delta, batch_id,
+                )
+
+        row = await self._with_disk_full_recovery(_do)
+        if row is None:
+            return None
+        raw = row["data"]
+        return json.loads(raw) if isinstance(raw, str) else dict(raw)
+
+    # ── Pruning ─────────────────────────────────────────────────
+
+    async def prune(
+        self,
+        status: str | None = None,
+        tag: str | None = None,
+        older_than_seconds: float | None = None,
+        queue: str | None = None,
+    ) -> int:
+        conditions = []
+        params: list[Any] = []
+        idx = 1
+
+        if status:
+            conditions.append(f"status=${idx}")
+            params.append(status)
+            idx += 1
+        if tag:
+            conditions.append(f"${idx} = ANY(tags)")
+            params.append(tag)
+            idx += 1
+        if older_than_seconds:
+            cutoff = _now_ts() - older_than_seconds
+            conditions.append(f"updated_at < ${idx}")
+            params.append(cutoff)
+            idx += 1
+        if queue:
+            conditions.append(f"queue=${idx}")
+            params.append(queue)
+            idx += 1
+
+        if not conditions:
+            return 0
+
+        where = " AND ".join(conditions)
+
+        async def _do():
+            async with self._pool.acquire() as conn:
+                return await conn.execute(f"DELETE FROM {self._jobs_table} WHERE {where}", *params)
+
+        result = await self._with_disk_full_recovery(_do)
+        return int(result.split()[-1])
+
+    async def flush(self, queue: str | None = None) -> None:
+        async def _do():
+            async with self._pool.acquire() as conn:
+                if queue:
+                    await conn.execute(f"DELETE FROM {self._jobs_table} WHERE queue=$1", queue)
+                else:
+                    await conn.execute(
+                        f"TRUNCATE {self._jobs_table}, {self._batches_table}, {self._metrics_table}, {self._supervisors_table}"
+                    )
+
+        await self._with_disk_full_recovery(_do)
+
+    async def prune_metrics(self, older_than_seconds: float) -> int:
+        cutoff = _now_ts() - older_than_seconds
+
+        async def _do():
+            async with self._pool.acquire() as conn:
+                return await conn.execute(
+                    f"DELETE FROM {self._metrics_table} WHERE recorded_at < $1", cutoff,
+                )
+
+        result = await self._with_disk_full_recovery(_do)
+        return int(result.split()[-1])
+
+    async def recent_throughput(
+        self, seconds: int = 60, queue: str | None = None,
+    ) -> dict[str, int]:
+        cutoff = _now_ts() - seconds
+        async with self._pool.acquire() as conn:
+            if queue:
+                rows = await conn.fetch(
+                    f"""SELECT metric, COUNT(*) AS cnt FROM {self._metrics_table}
+                        WHERE recorded_at > $1 AND queue = $2
+                        GROUP BY metric""",
+                    cutoff, queue,
+                )
+            else:
+                rows = await conn.fetch(
+                    f"""SELECT metric, COUNT(*) AS cnt FROM {self._metrics_table}
+                        WHERE recorded_at > $1
+                        GROUP BY metric""",
+                    cutoff,
+                )
+        out = {"processing": 0, "completed": 0, "failed": 0}
+        for r in rows:
+            if r["metric"] in out:
+                out[r["metric"]] = int(r["cnt"])
+        return out