89 lines
3.2 KiB
Python
89 lines
3.2 KiB
Python
import aiosqlite
|
|
|
|
from utils.db.models import Heartbeat, HeartbeatStatus
|
|
from utils.format import now
|
|
|
|
_COLUMNS = (
|
|
"id, name, status, text, period, created_at, "
|
|
"last_ping_at, last_ok_at, is_overdue, alerted_at"
|
|
)
|
|
|
|
|
|
def _row_to_heartbeat(row: aiosqlite.Row) -> Heartbeat:
|
|
return Heartbeat(
|
|
id=row["id"],
|
|
name=row["name"],
|
|
status=HeartbeatStatus(row["status"]),
|
|
text=row["text"],
|
|
period=row["period"],
|
|
created_at=row["created_at"],
|
|
last_ping_at=row["last_ping_at"],
|
|
last_ok_at=row["last_ok_at"],
|
|
is_overdue=bool(row["is_overdue"]),
|
|
alerted_at=row["alerted_at"],
|
|
)
|
|
|
|
|
|
class HeartbeatRepository:
|
|
def __init__(self, conn: aiosqlite.Connection) -> None:
|
|
self._conn = conn
|
|
|
|
async def all(self) -> list[Heartbeat]:
|
|
cursor = await self._conn.execute(
|
|
f"SELECT {_COLUMNS} FROM heartbeats ORDER BY name" # noqa: S608
|
|
)
|
|
return [_row_to_heartbeat(row) for row in await cursor.fetchall()]
|
|
|
|
async def get(self, heartbeat_id: int) -> Heartbeat | None:
|
|
cursor = await self._conn.execute(
|
|
f"SELECT {_COLUMNS} FROM heartbeats WHERE id = ?", # noqa: S608
|
|
(heartbeat_id,),
|
|
)
|
|
row = await cursor.fetchone()
|
|
return _row_to_heartbeat(row) if row else None
|
|
|
|
async def by_name(self, name: str) -> Heartbeat | None:
|
|
cursor = await self._conn.execute(
|
|
f"SELECT {_COLUMNS} FROM heartbeats WHERE name = ?", # noqa: S608
|
|
(name,),
|
|
)
|
|
row = await cursor.fetchone()
|
|
return _row_to_heartbeat(row) if row else None
|
|
|
|
async def record(
|
|
self, name: str, *, status: HeartbeatStatus, text: str, period: int
|
|
) -> Heartbeat:
|
|
stamp = now().isoformat()
|
|
ok_at = stamp if status is HeartbeatStatus.OK else None
|
|
await self._conn.execute(
|
|
"INSERT INTO heartbeats "
|
|
"(name, status, text, period, created_at, last_ping_at, last_ok_at) "
|
|
"VALUES (?, ?, ?, ?, ?, ?, ?) "
|
|
"ON CONFLICT(name) DO UPDATE SET status = excluded.status, "
|
|
"text = excluded.text, period = excluded.period, "
|
|
"last_ping_at = excluded.last_ping_at, "
|
|
"last_ok_at = COALESCE(excluded.last_ok_at, last_ok_at), "
|
|
"is_overdue = CASE WHEN excluded.last_ok_at IS NULL "
|
|
"THEN is_overdue ELSE 0 END, "
|
|
"alerted_at = CASE WHEN excluded.last_ok_at IS NULL "
|
|
"THEN alerted_at ELSE NULL END",
|
|
(name, status.value, text, period, stamp, stamp, ok_at),
|
|
)
|
|
await self._conn.commit()
|
|
heartbeat = await self.by_name(name)
|
|
assert heartbeat is not None
|
|
return heartbeat
|
|
|
|
async def mark_overdue(self, heartbeat_id: int, last_ping_at: str) -> bool:
|
|
cursor = await self._conn.execute(
|
|
"UPDATE heartbeats SET is_overdue = 1, alerted_at = ? "
|
|
"WHERE id = ? AND last_ping_at = ?",
|
|
(now().isoformat(), heartbeat_id, last_ping_at),
|
|
)
|
|
await self._conn.commit()
|
|
return cursor.rowcount > 0
|
|
|
|
async def delete(self, heartbeat_id: int) -> None:
|
|
await self._conn.execute("DELETE FROM heartbeats WHERE id = ?", (heartbeat_id,))
|
|
await self._conn.commit()
|