127 lines
4 KiB
Python
127 lines
4 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import aiosqlite
|
||
|
|
|
||
|
|
_SCHEMA = """
|
||
|
|
CREATE TABLE IF NOT EXISTS entity_baseline (
|
||
|
|
entity_id TEXT PRIMARY KEY,
|
||
|
|
state TEXT NOT NULL,
|
||
|
|
attributes TEXT NOT NULL DEFAULT '{}',
|
||
|
|
updated_at REAL NOT NULL
|
||
|
|
);
|
||
|
|
|
||
|
|
CREATE TABLE IF NOT EXISTS check_history (
|
||
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
|
|
check_name TEXT NOT NULL,
|
||
|
|
ran_at REAL NOT NULL,
|
||
|
|
healthy INTEGER NOT NULL,
|
||
|
|
message TEXT NOT NULL DEFAULT '',
|
||
|
|
payload TEXT NOT NULL DEFAULT '{}'
|
||
|
|
);
|
||
|
|
|
||
|
|
CREATE TABLE IF NOT EXISTS alerts_sent (
|
||
|
|
alert_key TEXT PRIMARY KEY,
|
||
|
|
sent_at REAL NOT NULL
|
||
|
|
);
|
||
|
|
"""
|
||
|
|
|
||
|
|
|
||
|
|
class Storage:
|
||
|
|
def __init__(self, db_path: Path) -> None:
|
||
|
|
self._db_path = db_path
|
||
|
|
self._db: aiosqlite.Connection | None = None
|
||
|
|
|
||
|
|
async def open(self) -> None:
|
||
|
|
self._db_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
self._db = await aiosqlite.connect(self._db_path)
|
||
|
|
self._db.row_factory = aiosqlite.Row
|
||
|
|
await self._db.executescript(_SCHEMA)
|
||
|
|
await self._db.commit()
|
||
|
|
|
||
|
|
async def close(self) -> None:
|
||
|
|
if self._db:
|
||
|
|
await self._db.close()
|
||
|
|
self._db = None
|
||
|
|
|
||
|
|
def _conn(self) -> aiosqlite.Connection:
|
||
|
|
if self._db is None:
|
||
|
|
raise RuntimeError("Storage not open — call await storage.open() first")
|
||
|
|
return self._db
|
||
|
|
|
||
|
|
# ------------------------------------------------------------------
|
||
|
|
# entity_baseline
|
||
|
|
# ------------------------------------------------------------------
|
||
|
|
|
||
|
|
async def upsert_entity_baseline(
|
||
|
|
self, entity_id: str, state: str, attributes: str, updated_at: float
|
||
|
|
) -> None:
|
||
|
|
await self._conn().execute(
|
||
|
|
"""
|
||
|
|
INSERT INTO entity_baseline (entity_id, state, attributes, updated_at)
|
||
|
|
VALUES (?, ?, ?, ?)
|
||
|
|
ON CONFLICT(entity_id) DO UPDATE SET
|
||
|
|
state = excluded.state,
|
||
|
|
attributes = excluded.attributes,
|
||
|
|
updated_at = excluded.updated_at
|
||
|
|
""",
|
||
|
|
(entity_id, state, attributes, updated_at),
|
||
|
|
)
|
||
|
|
await self._conn().commit()
|
||
|
|
|
||
|
|
async def get_entity_baseline(self, entity_id: str) -> dict[str, Any] | None:
|
||
|
|
async with self._conn().execute(
|
||
|
|
"SELECT * FROM entity_baseline WHERE entity_id = ?", (entity_id,)
|
||
|
|
) as cur:
|
||
|
|
row = await cur.fetchone()
|
||
|
|
return dict(row) if row else None
|
||
|
|
|
||
|
|
# ------------------------------------------------------------------
|
||
|
|
# check_history
|
||
|
|
# ------------------------------------------------------------------
|
||
|
|
|
||
|
|
async def record_check(
|
||
|
|
self,
|
||
|
|
check_name: str,
|
||
|
|
ran_at: float,
|
||
|
|
healthy: bool,
|
||
|
|
message: str,
|
||
|
|
payload: str,
|
||
|
|
) -> None:
|
||
|
|
await self._conn().execute(
|
||
|
|
"""
|
||
|
|
INSERT INTO check_history (check_name, ran_at, healthy, message, payload)
|
||
|
|
VALUES (?, ?, ?, ?, ?)
|
||
|
|
""",
|
||
|
|
(check_name, ran_at, int(healthy), message, payload),
|
||
|
|
)
|
||
|
|
await self._conn().commit()
|
||
|
|
|
||
|
|
# ------------------------------------------------------------------
|
||
|
|
# alerts_sent (dedup gate)
|
||
|
|
# ------------------------------------------------------------------
|
||
|
|
|
||
|
|
async def was_alert_sent(self, alert_key: str, within_seconds: float) -> bool:
|
||
|
|
import time
|
||
|
|
|
||
|
|
cutoff = time.time() - within_seconds
|
||
|
|
async with self._conn().execute(
|
||
|
|
"SELECT sent_at FROM alerts_sent WHERE alert_key = ? AND sent_at > ?",
|
||
|
|
(alert_key, cutoff),
|
||
|
|
) as cur:
|
||
|
|
return (await cur.fetchone()) is not None
|
||
|
|
|
||
|
|
async def mark_alert_sent(self, alert_key: str) -> None:
|
||
|
|
import time
|
||
|
|
|
||
|
|
await self._conn().execute(
|
||
|
|
"""
|
||
|
|
INSERT INTO alerts_sent (alert_key, sent_at) VALUES (?, ?)
|
||
|
|
ON CONFLICT(alert_key) DO UPDATE SET sent_at = excluded.sent_at
|
||
|
|
""",
|
||
|
|
(alert_key, time.time()),
|
||
|
|
)
|
||
|
|
await self._conn().commit()
|