Prepare KarlDNS production deployment

karldns/store/scheduler.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from datetime import UTC, datetime, timedelta
+
+from .types import ScheduledTaskRow
+
+
+def _now() -> datetime:
+    return datetime.now(UTC)
+
+
+def _iso(value: datetime) -> str:
+    return value.isoformat()
+
+
+class SchedulerMixin:
+    def ensure_scheduled_task(self, name: str, interval_seconds: int) -> None:
+        now = _now()
+        with self.connect() as conn:
+            conn.execute(
+                """
+                INSERT INTO scheduled_tasks(name, interval_seconds, next_run_at, updated_at)
+                VALUES (?, ?, ?, ?)
+                ON CONFLICT(name) DO UPDATE SET
+                  interval_seconds=excluded.interval_seconds,
+                  updated_at=excluded.updated_at
+                """,
+                (name, interval_seconds, _iso(now), _iso(now)),
+            )
+
+    def acquire_due_task(self, name: str, lease_owner: str, lease_seconds: int) -> bool:
+        now = _now()
+        lease_expires_at = now + timedelta(seconds=lease_seconds)
+        with self.connect() as conn:
+            cursor = conn.execute(
+                """
+                UPDATE scheduled_tasks
+                SET lease_owner = ?,
+                    lease_expires_at = ?,
+                    last_started_at = ?,
+                    last_status = 'running',
+                    last_error = '',
+                    updated_at = ?
+                WHERE name = ?
+                  AND next_run_at <= ?
+                  AND (lease_expires_at IS NULL OR lease_expires_at <= ?)
+                """,
+                (
+                    lease_owner,
+                    _iso(lease_expires_at),
+                    _iso(now),
+                    _iso(now),
+                    name,
+                    _iso(now),
+                    _iso(now),
+                ),
+            )
+        return cursor.rowcount == 1
+
+    def complete_scheduled_task(
+        self,
+        name: str,
+        interval_seconds: int,
+        success: bool,
+        error: str = "",
+    ) -> None:
+        now = _now()
+        next_run_at = now + timedelta(seconds=interval_seconds)
+        status = "ok" if success else "error"
+        with self.connect() as conn:
+            conn.execute(
+                """
+                UPDATE scheduled_tasks
+                SET next_run_at = ?,
+                    lease_owner = NULL,
+                    lease_expires_at = NULL,
+                    last_finished_at = ?,
+                    last_status = ?,
+                    last_error = ?,
+                    interval_seconds = ?,
+                    updated_at = ?
+                WHERE name = ?
+                """,
+                (_iso(next_run_at), _iso(now), status, error[:500], interval_seconds, _iso(now), name),
+            )
+
+    def list_scheduled_tasks(self) -> list[ScheduledTaskRow]:
+        with self.connect() as conn:
+            rows = conn.execute(
+                """
+                SELECT name, interval_seconds, next_run_at, lease_owner, lease_expires_at,
+                       last_started_at, last_finished_at, last_status, last_error, updated_at
+                FROM scheduled_tasks
+                ORDER BY name
+                """
+            ).fetchall()
+        return [ScheduledTaskRow(**dict(row)) for row in rows]