|
# retoor <retoor@molodetz.nl>
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from datetime import datetime, timezone
|
|
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]
|
|
if not table.has_column("deleted_at"):
|
|
table.create_column_by_example("deleted_at", "")
|
|
if not table.has_column("deleted_by"):
|
|
table.create_column_by_example("deleted_by", "")
|
|
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(
|
|
{"deleted_at": None, "deleted_by": None, **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, deleted_at=None, **self._scope)
|
|
|
|
def list(
|
|
self, enabled_only: bool = False, status: str | None = None
|
|
) -> list[dict[str, Any]]:
|
|
criteria: dict[str, Any] = dict(self._scope)
|
|
criteria["deleted_at"] = None
|
|
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:
|
|
row = self._table.find_one(uid=uid, deleted_at=None, **self._scope)
|
|
if not row:
|
|
return False
|
|
self._table.update(
|
|
{
|
|
"uid": uid,
|
|
"deleted_at": datetime.now(timezone.utc).isoformat(),
|
|
"deleted_by": f"{self._owner_kind}:{self._owner_id}",
|
|
**self._scope,
|
|
},
|
|
["uid", "owner_kind", "owner_id"],
|
|
)
|
|
logger.info("Task soft-deleted uid=%s", uid)
|
|
return True
|
|
|
|
def recover_running(self) -> int:
|
|
if TABLE not in self._db.tables:
|
|
return 0
|
|
stuck = list(self._table.find(status="running", deleted_at=None, **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", deleted_at=None, **self._scope
|
|
)
|
|
return [
|
|
row
|
|
for row in rows
|
|
if row.get("next_run_at") and row["next_run_at"] <= now_iso
|
|
]
|