From 69e7ad79ac1f77e285c5727e6fbfe2270451ce3b Mon Sep 17 00:00:00 2001 From: h Date: Wed, 14 Oct 2026 05:20:53 +0200 Subject: [PATCH] feat(heartbeats): ping endpoint with silence alerts, menu screen and digest block --- .env.example | 8 ++ README.md | 4 + src/healthbot/__init__.py | 3 + src/healthbot/digest.py | 13 +++ src/healthbot/handlers/__init__.py | 3 +- src/healthbot/handlers/heartbeats.py | 127 +++++++++++++++++++++++++ src/healthbot/heartbeats.py | 85 +++++++++++++++++ src/healthbot/keyboards/__init__.py | 53 ++++++++--- src/healthbot/web.py | 53 +++++++++++ src/utils/db/__init__.py | 18 +++- src/utils/db/models.py | 27 ++++++ src/utils/db/repositories/__init__.py | 2 + src/utils/db/repositories/heartbeat.py | 88 +++++++++++++++++ src/utils/env.py | 6 ++ 14 files changed, 473 insertions(+), 17 deletions(-) create mode 100644 src/healthbot/handlers/heartbeats.py create mode 100644 src/healthbot/heartbeats.py create mode 100644 src/utils/db/repositories/heartbeat.py diff --git a/.env.example b/.env.example index a6afb21..df533b5 100644 --- a/.env.example +++ b/.env.example @@ -35,6 +35,14 @@ MONITOR__BODY_PREVIEW=400 ## же можно дописать прямо в URL: .../komodo/alert?token=... # WEB__TOKEN= +## --- Отметки от заданий (необязательно) --------------------------------- +## Задание после каждого запуска шлёт POST /ping/<имя> с заголовком +## X-Ping-Token и телом {"status": "ok" | "fail", "text": "...", "period": 93600}. +## Нет "ok" дольше period секунд: бот сам поднимает тревогу. +## Работает на том же приёмнике (WEB__ENABLED=true), но этот адрес должен быть +## доступен заданию снаружи. Пустой токен выключает /ping целиком. +# WEB__PING_TOKEN= + DB__PATH=data/healthbot.db LOG__LEVEL=INFO diff --git a/README.md b/README.md index ce62ff0..4ee2745 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,7 @@ # healthbot [![AI Slop Inside](https://sladge.net/badge.svg)](https://sladge.net) + +## Heartbeats + +Jobs report with `POST /ping/{name}`, header `X-Ping-Token: $WEB__PING_TOKEN`, body `{"status": "ok" | "fail", "text": "...", "period": 93600}`. No `ok` for longer than `period` seconds and the bot raises the alarm. diff --git a/src/healthbot/__init__.py b/src/healthbot/__init__.py index 5a2125b..41b2d2d 100644 --- a/src/healthbot/__init__.py +++ b/src/healthbot/__init__.py @@ -13,6 +13,7 @@ async def runner() -> None: from . import handlers # noqa: PLC0415 from .common import bot, dp # noqa: PLC0415 from .digest import run_digest # noqa: PLC0415 + from .heartbeats import run_heartbeats # noqa: PLC0415 from .komodo import komodo # noqa: PLC0415 from .monitor import run_monitor # noqa: PLC0415 from .web import run_web # noqa: PLC0415 @@ -22,6 +23,8 @@ async def runner() -> None: await bot.delete_webhook(drop_pending_updates=True) background = [asyncio.create_task(run_monitor()), asyncio.create_task(run_digest())] + if env.web.pings: + background.append(asyncio.create_task(run_heartbeats())) web = await run_web() if env.web.enabled else None try: await dp.start_polling(bot) diff --git a/src/healthbot/digest.py b/src/healthbot/digest.py index 01f160a..4e89d3e 100644 --- a/src/healthbot/digest.py +++ b/src/healthbot/digest.py @@ -1,6 +1,7 @@ import asyncio from datetime import timedelta +from healthbot.heartbeats import last_ok from healthbot.keyboards import back_to_menu from healthbot.komodo import KomodoError, Update, komodo from healthbot.notify import notify @@ -89,10 +90,22 @@ async def _fleet_block() -> list[str]: return lines +async def _heartbeats_block() -> list[str]: + heartbeats = await db.heartbeats.all() + if not heartbeats: + return [] + good = sum(h.healthy for h in heartbeats) + lines = [f"💓 Отметки: в порядке {good} из {len(heartbeats)}"] + lines.extend(f"{h.mark} {esc(h.name)}: {last_ok(h)}" for h in heartbeats) + return lines + + async def digest_text() -> str: blocks = [await _services_block()] if env.komodo.enabled: blocks.append(await _fleet_block()) + if env.web.pings: + blocks.append(await _heartbeats_block()) lines = ["🗞 Сводка за сутки"] for block in blocks: if block: diff --git a/src/healthbot/handlers/__init__.py b/src/healthbot/handlers/__init__.py index 33d036d..25e07d9 100644 --- a/src/healthbot/handlers/__init__.py +++ b/src/healthbot/handlers/__init__.py @@ -2,12 +2,13 @@ from aiogram import Router, types from utils.logging import logger -from . import actions, add, fleet, menu, runs, services, settings +from . import actions, add, fleet, heartbeats, menu, runs, services, settings router = Router() router.include_routers( menu.router, services.router, + heartbeats.router, add.router, fleet.router, actions.router, diff --git a/src/healthbot/handlers/heartbeats.py b/src/healthbot/handlers/heartbeats.py new file mode 100644 index 0000000..aa946db --- /dev/null +++ b/src/healthbot/handlers/heartbeats.py @@ -0,0 +1,127 @@ +from aiogram import F, Router, types +from aiogram.fsm.context import FSMContext + +from healthbot.filters import Admin +from healthbot.heartbeats import last_ok, silence +from healthbot.keyboards import ( + back_to_menu, + confirm_delete, + heartbeat_card, + heartbeats_list, +) +from healthbot.ui import esc, int_arg, pages, show +from utils.db import db +from utils.db.models import Heartbeat, HeartbeatStatus +from utils.format import ago, human_duration, parse, short, stamp + +router = Router() +router.callback_query.filter(Admin()) + +PAGE = 10 +SHOWN_TEXT = 80 + + +async def list_view(page: int = 0) -> tuple[str, types.InlineKeyboardMarkup]: + heartbeats = sorted(await db.heartbeats.all(), key=lambda h: h.healthy) + if not heartbeats: + return ( + "💓 Отметки\n" + "Пока никто не отчитывался. Первый запрос на " + "/ping/имя заведёт отметку сам.", + back_to_menu(), + ) + total = pages(len(heartbeats), PAGE) + page = min(page, total - 1) + chunk = heartbeats[page * PAGE : (page + 1) * PAGE] + good = sum(h.healthy for h in heartbeats) + lines = [f"💓 Отметки 🟢 {good} 🔴 {len(heartbeats) - good}", ""] + for heartbeat in chunk: + lines.append( + f"{heartbeat.mark} {esc(heartbeat.name)} · {last_ok(heartbeat)}" + ) + if heartbeat.text: + lines.append(esc(short(heartbeat.text, SHOWN_TEXT))) + return "\n".join(lines), heartbeats_list(chunk, page, total) + + +@router.callback_query(F.data.startswith("hb:list:")) +async def hb_list(callback: types.CallbackQuery, state: FSMContext) -> None: + await state.clear() + text, markup = await list_view(int_arg(callback.data)) + await show(callback, text, markup) + await callback.answer() + + +def _state(heartbeat: Heartbeat) -> str: + if heartbeat.is_overdue: + return f"молчит {human_duration(silence(heartbeat))}" + if heartbeat.status is HeartbeatStatus.FAIL: + return "сбой" + return "в порядке" + + +def _moment(value: str | None, tz: str) -> str: + if value is None: + return "ещё не было" + return f"{stamp(parse(value), tz)}, {ago(parse(value))}" + + +def card_view(heartbeat: Heartbeat) -> str: + tz = db.settings.get("tz") + lines = [ + f"{heartbeat.mark} {esc(heartbeat.name)}", + "", + f"Статус: {_state(heartbeat)}", + f"Период: раз в {human_duration(heartbeat.period)}", + f"Последний успех: {_moment(heartbeat.last_ok_at, tz)}", + f"Последняя отметка: {_moment(heartbeat.last_ping_at, tz)}", + ] + if heartbeat.text: + lines.append(f"
{esc(heartbeat.text)}
") + return "\n".join(lines) + + +async def _load(callback: types.CallbackQuery) -> Heartbeat | None: + heartbeat = await db.heartbeats.get(int_arg(callback.data)) + if heartbeat is None: + await callback.answer("Отметка не найдена", show_alert=True) + return heartbeat + + +@router.callback_query(F.data.startswith("hb:open:")) +async def hb_open(callback: types.CallbackQuery) -> None: + heartbeat = await _load(callback) + if heartbeat is None: + return + await show(callback, card_view(heartbeat), heartbeat_card(heartbeat.id)) + await callback.answer() + + +@router.callback_query(F.data.startswith("go:hb:")) +async def go_heartbeat(callback: types.CallbackQuery) -> None: + heartbeat = await _load(callback) + if heartbeat is None: + return + await show(callback, card_view(heartbeat), heartbeat_card(heartbeat.id), new=True) + await callback.answer() + + +@router.callback_query(F.data.startswith("hb:del:")) +async def hb_del(callback: types.CallbackQuery) -> None: + heartbeat = await _load(callback) + if heartbeat is None: + return + await show( + callback, + f"Удалить {esc(heartbeat.name)}? Новая отметка заведёт его заново.", + confirm_delete(heartbeat.id, "hb"), + ) + await callback.answer() + + +@router.callback_query(F.data.startswith("hb:delyes:")) +async def hb_del_yes(callback: types.CallbackQuery) -> None: + await db.heartbeats.delete(int_arg(callback.data)) + text, markup = await list_view() + await show(callback, text, markup) + await callback.answer() diff --git a/src/healthbot/heartbeats.py b/src/healthbot/heartbeats.py new file mode 100644 index 0000000..08badbe --- /dev/null +++ b/src/healthbot/heartbeats.py @@ -0,0 +1,85 @@ +import asyncio + +from aiogram.types import InlineKeyboardMarkup +from aiogram.utils.keyboard import InlineKeyboardBuilder + +from healthbot.notify import notify +from healthbot.ui import esc +from utils.db import db +from utils.db.models import Heartbeat, HeartbeatStatus +from utils.format import ago, human_duration, now, parse +from utils.logging import logger + +TICK = 60 +REMIND_EVERY = 86400 +DEFAULT_PERIOD = 93600 +MIN_PERIOD = 300 +MAX_PERIOD = 2678400 +MAX_TEXT = 1000 + + +def silence(heartbeat: Heartbeat) -> float: + since = parse(heartbeat.last_ok_at or heartbeat.created_at) + return (now() - since).total_seconds() + + +def last_ok(heartbeat: Heartbeat) -> str: + if heartbeat.last_ok_at is None: + return "успеха ещё не было" + return f"успех {ago(parse(heartbeat.last_ok_at))}" + + +def _open_kb(heartbeat: Heartbeat) -> InlineKeyboardMarkup: + kb = InlineKeyboardBuilder() + kb.button(text="📂 Открыть", callback_data=f"go:hb:{heartbeat.id}") + return kb.as_markup() + + +def _message(heartbeat: Heartbeat, icon: str, tail: str) -> str: + head = f"{icon} {esc(heartbeat.name)} {tail}" + if not heartbeat.text: + return head + return f"{head}\n
{esc(heartbeat.text)}
" + + +async def on_ping(name: str, status: HeartbeatStatus, text: str, period: int) -> None: + previous = await db.heartbeats.by_name(name) + heartbeat = await db.heartbeats.record( + name, status=status, text=text, period=period + ) + if previous is None: + logger.info(f"Heartbeat registered: {name}") + if status is HeartbeatStatus.FAIL: + await notify(_message(heartbeat, "🔴", "сбой"), _open_kb(heartbeat)) + logger.warning(f"[red]FAIL[/] {name}") + elif previous is not None and not previous.healthy: + await notify(_message(heartbeat, "✅", "снова отчитался"), _open_kb(heartbeat)) + logger.info(f"[green]OK[/] {name}") + + +async def check_overdue() -> None: + for heartbeat in await db.heartbeats.all(): + idle = silence(heartbeat) + if idle <= heartbeat.period: + continue + if heartbeat.alerted_at is not None: + since = (now() - parse(heartbeat.alerted_at)).total_seconds() + if since < REMIND_EVERY: + continue + if not await db.heartbeats.mark_overdue(heartbeat.id, heartbeat.last_ping_at): + continue + tail = ( + f"молчит {human_duration(idle)}, " + f"ждал отметку раз в {human_duration(heartbeat.period)}" + ) + await notify(_message(heartbeat, "🔴", tail), _open_kb(heartbeat)) + logger.warning(f"[red]SILENT[/] {heartbeat.name}") + + +async def run_heartbeats() -> None: + while True: + try: + await check_overdue() + except Exception: + logger.exception("Heartbeat check failed") + await asyncio.sleep(TICK) diff --git a/src/healthbot/keyboards/__init__.py b/src/healthbot/keyboards/__init__.py index befadf6..8841886 100644 --- a/src/healthbot/keyboards/__init__.py +++ b/src/healthbot/keyboards/__init__.py @@ -2,7 +2,7 @@ from aiogram.types import InlineKeyboardButton, InlineKeyboardMarkup from aiogram.utils.keyboard import InlineKeyboardBuilder from healthbot.ui import DANGER, PRIMARY, label, pager, pages, with_icon -from utils.db.models import Pin, PinMode, Service +from utils.db.models import Heartbeat, Pin, PinMode, Service from utils.env import env SERVICES_PAGE = 12 @@ -23,19 +23,17 @@ def main_menu(pins: list[Pin], *, actions: bool) -> InlineKeyboardMarkup: ) rows += [2] * (len(pins) // 2) + [1] * (len(pins) % 2) - kb.button(text="📋 Сервисы", callback_data="menu:list:0") + sections = [("📋 Сервисы", "menu:list:0")] if env.komodo.enabled: - kb.button(text="🛰 Флот", callback_data="fleet:main") + sections.append(("🛰 Флот", "fleet:main")) if actions: - kb.button(text="🎛 Действия", callback_data="act:list:a:0") - kb.button(text="📊 Статистика", callback_data="menu:stats") - rows += [2, 2] - else: - kb.button(text="📊 Статистика", callback_data="menu:stats") - rows += [2, 1] - else: - kb.button(text="📊 Статистика", callback_data="menu:stats") - rows += [2] + sections.append(("🎛 Действия", "act:list:a:0")) + if env.web.pings: + sections.append(("💓 Отметки", "hb:list:0")) + sections.append(("📊 Статистика", "menu:stats")) + for text, data in sections: + kb.button(text=text, callback_data=data) + rows += [2] * (len(sections) // 2) + [1] * (len(sections) % 2) kb.button(text="➕ Добавить сервис", callback_data="menu:add") kb.button(text="⚙️ Настройки", callback_data="set:main") kb.button(text="🔄 Обновить", callback_data="menu:main") @@ -91,10 +89,35 @@ def service_card(service: Service) -> InlineKeyboardMarkup: return kb.as_markup() -def confirm_delete(service_id: int) -> InlineKeyboardMarkup: +def confirm_delete(target_id: int, prefix: str = "svc") -> InlineKeyboardMarkup: kb = InlineKeyboardBuilder() - kb.button(text="🗑 Удалить", callback_data=f"svc:delyes:{service_id}", style=DANGER) - kb.button(text="✖️ Отмена", callback_data=f"svc:open:{service_id}") + kb.button( + text="🗑 Удалить", callback_data=f"{prefix}:delyes:{target_id}", style=DANGER + ) + kb.button(text="✖️ Отмена", callback_data=f"{prefix}:open:{target_id}") + kb.adjust(2) + return kb.as_markup() + + +def heartbeats_list( + chunk: list[Heartbeat], page: int, total: int +) -> InlineKeyboardMarkup: + kb = InlineKeyboardBuilder() + for heartbeat in chunk: + kb.button( + text=f"{heartbeat.mark} {heartbeat.name}", + callback_data=f"hb:open:{heartbeat.id}", + ) + nav = pager(kb, "hb:list", page, total) + kb.button(text="🏠 Меню", callback_data="menu:main") + kb.adjust(*[1] * len(chunk), *([nav] if nav else []), 1) + return kb.as_markup() + + +def heartbeat_card(heartbeat_id: int) -> InlineKeyboardMarkup: + kb = InlineKeyboardBuilder() + kb.button(text="🗑 Удалить", callback_data=f"hb:del:{heartbeat_id}") + kb.button(text="⬅️ К списку", callback_data="hb:list:0") kb.adjust(2) return kb.as_markup() diff --git a/src/healthbot/web.py b/src/healthbot/web.py index 9066327..eda0d4c 100644 --- a/src/healthbot/web.py +++ b/src/healthbot/web.py @@ -1,15 +1,27 @@ import hmac import json +import re from aiohttp import web from healthbot.alerts import alert_keyboard, alert_kind, format_alert +from healthbot.heartbeats import ( + DEFAULT_PERIOD, + MAX_PERIOD, + MAX_TEXT, + MIN_PERIOD, + on_ping, +) from healthbot.notify import notify from healthbot.ui import esc from utils.db import db +from utils.db.models import HeartbeatStatus from utils.env import env from utils.logging import logger +PING_NAME = re.compile(r"[a-z0-9][a-z0-9-]{0,39}") +PING_STATUSES = tuple(HeartbeatStatus) + def _authorized(request: web.Request) -> bool: expected = env.web.token.get_secret_value() @@ -52,6 +64,46 @@ async def handle_alert(request: web.Request) -> web.Response: return web.json_response({"ok": True}) +def _ping_fields(payload: object) -> tuple[HeartbeatStatus, str, int] | None: + if not isinstance(payload, dict): + return None + status = payload.get("status") + period = payload.get("period") + if period is None: + period = DEFAULT_PERIOD + if status not in PING_STATUSES or type(period) is not int: + return None + text = str(payload.get("text") or "").strip()[:MAX_TEXT] + return HeartbeatStatus(status), text, min(max(period, MIN_PERIOD), MAX_PERIOD) + + +async def handle_ping(request: web.Request) -> web.Response: + expected = env.web.ping_token.get_secret_value() + if not expected: + return web.json_response({"error": "not found"}, status=404) + given = request.headers.get("X-Ping-Token") or "" + if not hmac.compare_digest( + given.encode(errors="surrogateescape"), expected.encode() + ): + logger.warning(f"Ping without a valid token from {request.remote}") + return web.json_response({"error": "unauthorized"}, status=401) + + name = request.match_info["name"] + if not PING_NAME.fullmatch(name): + return web.json_response({"error": "bad name"}, status=400) + + try: + payload = json.loads(await request.text()) + except ValueError: + return web.json_response({"error": "bad json"}, status=400) + fields = _ping_fields(payload) + if fields is None: + return web.json_response({"error": "bad body"}, status=400) + + await on_ping(name, *fields) + return web.json_response({"ok": True}) + + async def handle_health(_: web.Request) -> web.Response: return web.json_response({"status": "ok"}) @@ -59,6 +111,7 @@ async def handle_health(_: web.Request) -> web.Response: async def run_web() -> web.AppRunner: app = web.Application() app.router.add_post("/komodo/alert", handle_alert) + app.router.add_post("/ping/{name}", handle_ping) app.router.add_get("/health", handle_health) runner = web.AppRunner(app) diff --git a/src/utils/db/__init__.py b/src/utils/db/__init__.py index 51430ee..5247224 100644 --- a/src/utils/db/__init__.py +++ b/src/utils/db/__init__.py @@ -7,6 +7,7 @@ from utils.env import env from utils.logging import logger from .repositories import ( + HeartbeatRepository, IncidentRepository, PinRepository, ServiceRepository, @@ -14,7 +15,7 @@ from .repositories import ( StatRepository, ) -SCHEMA_VERSION = 2 +SCHEMA_VERSION = 3 SCHEMA = """ CREATE TABLE IF NOT EXISTS services ( @@ -67,6 +68,19 @@ CREATE TABLE IF NOT EXISTS hourly ( latency_max INTEGER NOT NULL DEFAULT 0, PRIMARY KEY (service_id, hour) ); + +CREATE TABLE IF NOT EXISTS heartbeats ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + name TEXT NOT NULL UNIQUE, + status TEXT NOT NULL, + text TEXT NOT NULL DEFAULT '', + period INTEGER NOT NULL, + created_at TEXT NOT NULL, + last_ping_at TEXT NOT NULL, + last_ok_at TEXT, + is_overdue INTEGER NOT NULL DEFAULT 0, + alerted_at TEXT +); """ MIGRATIONS: tuple[tuple[str, str, str], ...] = ( @@ -102,6 +116,7 @@ class Database: self._path = path self._conn: aiosqlite.Connection | None = None self.services: ServiceRepository = None # type: ignore[assignment] + self.heartbeats: HeartbeatRepository = None # type: ignore[assignment] self.incidents: IncidentRepository = None # type: ignore[assignment] self.pins: PinRepository = None # type: ignore[assignment] self.settings: SettingRepository = None # type: ignore[assignment] @@ -118,6 +133,7 @@ class Database: await self._conn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}") await self._conn.commit() self.services = ServiceRepository(self._conn) + self.heartbeats = HeartbeatRepository(self._conn) self.incidents = IncidentRepository(self._conn) self.pins = PinRepository(self._conn) self.settings = SettingRepository(self._conn) diff --git a/src/utils/db/models.py b/src/utils/db/models.py index 5ab7a5c..af6028c 100644 --- a/src/utils/db/models.py +++ b/src/utils/db/models.py @@ -15,6 +15,11 @@ class PinMode(StrEnum): FORM = "form" +class HeartbeatStatus(StrEnum): + OK = "ok" + FAIL = "fail" + + @dataclass(slots=True) class Service: id: int @@ -79,6 +84,28 @@ class Pin: position: int = 0 +@dataclass(slots=True) +class Heartbeat: + id: int + name: str + status: HeartbeatStatus + text: str + period: int + created_at: str + last_ping_at: str + last_ok_at: str | None = None + is_overdue: bool = False + alerted_at: str | None = None + + @property + def healthy(self) -> bool: + return self.status is HeartbeatStatus.OK and not self.is_overdue + + @property + def mark(self) -> str: + return "🟢" if self.healthy else "🔴" + + @dataclass(slots=True) class Hour: hour: int diff --git a/src/utils/db/repositories/__init__.py b/src/utils/db/repositories/__init__.py index 2458f3d..2546500 100644 --- a/src/utils/db/repositories/__init__.py +++ b/src/utils/db/repositories/__init__.py @@ -1,3 +1,4 @@ +from .heartbeat import HeartbeatRepository from .incident import IncidentRepository from .pin import PinRepository from .service import ServiceRepository @@ -5,6 +6,7 @@ from .setting import SettingRepository from .stat import StatRepository __all__ = [ + "HeartbeatRepository", "IncidentRepository", "PinRepository", "ServiceRepository", diff --git a/src/utils/db/repositories/heartbeat.py b/src/utils/db/repositories/heartbeat.py new file mode 100644 index 0000000..f3ad0eb --- /dev/null +++ b/src/utils/db/repositories/heartbeat.py @@ -0,0 +1,88 @@ +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() diff --git a/src/utils/env.py b/src/utils/env.py index 68b4a20..ca3409e 100644 --- a/src/utils/env.py +++ b/src/utils/env.py @@ -37,6 +37,12 @@ class WebSettings(BaseSettings): host: str = "0.0.0.0" # noqa: S104 port: int = 8080 token: SecretStr = SecretStr("") + ping_token: SecretStr = SecretStr("") + + @computed_field + @property + def pings(self) -> bool: + return bool(self.enabled and self.ping_token.get_secret_value()) class DbSettings(BaseSettings):