# retoor from __future__ import annotations import logging from typing import Any import dataset logger = logging.getLogger("devii.tasks.store") TABLE = "devii_tasks" INDEXED_COLUMNS = (["owner_kind", "owner_id"], ["uid"], ["enabled"], ["next_run_at"], ["status"]) def memory_db() -> Any: return dataset.connect("sqlite:///:memory:") class TaskStore: def __init__(self, db: Any, owner_kind: str, owner_id: str) -> None: self._db = db self._owner_kind = owner_kind self._owner_id = owner_id self._ensure_indexes() def _ensure_indexes(self) -> None: if TABLE not in self._db.tables: return table = self._db[TABLE] for columns in INDEXED_COLUMNS: table.create_index(columns) @property def _table(self) -> Any: return self._db[TABLE] @property def _scope(self) -> dict[str, str]: return {"owner_kind": self._owner_kind, "owner_id": self._owner_id} def create(self, record: dict[str, Any]) -> None: self._table.insert({**record, **self._scope}) logger.info("Task created uid=%s owner=%s/%s", record.get("uid"), self._owner_kind, self._owner_id) def get(self, uid: str) -> dict[str, Any] | None: return self._table.find_one(uid=uid, **self._scope) def list(self, enabled_only: bool = False, status: str | None = None) -> list[dict[str, Any]]: criteria: dict[str, Any] = dict(self._scope) if enabled_only: criteria["enabled"] = True if status: criteria["status"] = status return list(self._table.find(**criteria)) def update(self, uid: str, changes: dict[str, Any]) -> None: changes = {**changes, "uid": uid, **self._scope} self._table.update(changes, ["uid", "owner_kind", "owner_id"]) logger.debug("Task updated uid=%s changes=%s", uid, list(changes)) def delete(self, uid: str) -> bool: deleted = self._table.delete(uid=uid, **self._scope) logger.info("Task deleted uid=%s ok=%s", uid, deleted) return bool(deleted) def recover_running(self) -> int: if TABLE not in self._db.tables: return 0 stuck = list(self._table.find(status="running", **self._scope)) for row in stuck: self._table.update({"uid": row["uid"], "status": "pending", **self._scope}, ["uid", "owner_kind", "owner_id"]) if stuck: logger.info("Recovered %d task(s) stuck in running", len(stuck)) return len(stuck) def due(self, now_iso: str) -> list[dict[str, Any]]: rows = self._table.find(enabled=True, status="pending", **self._scope) return [row for row in rows if row.get("next_run_at") and row["next_run_at"] <= now_iso]