135 lines
4.0 KiB
Python
135 lines
4.0 KiB
Python
import json
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
|
|
import asyncpg
|
|
|
|
from utils.read.models import AlertView, Page, WatchView
|
|
|
|
WATCHES_CHANGED_CHANNEL = "watches_changed"
|
|
ALERTS_CHANGED_CHANNEL = "alerts_changed"
|
|
|
|
_WATCH_COLS = "id, account_id, kind, params, enabled, created_at, updated_at"
|
|
_ALERT_COLS = "id, account_id, watch_id, ts, payload, seen, created_at"
|
|
|
|
|
|
def _to_watch(row: asyncpg.Record) -> WatchView:
|
|
data = dict(row)
|
|
data["params"] = json.loads(data["params"])
|
|
return WatchView(**data)
|
|
|
|
|
|
def _to_alert(row: asyncpg.Record) -> AlertView:
|
|
data = dict(row)
|
|
data["payload"] = json.loads(data["payload"])
|
|
return AlertView(**data)
|
|
|
|
|
|
async def list_watches(pool: asyncpg.Pool, account_id: int) -> list[WatchView]:
|
|
rows = await pool.fetch(
|
|
f"SELECT {_WATCH_COLS} FROM watches WHERE account_id = $1 " # noqa: S608
|
|
"ORDER BY id DESC",
|
|
account_id,
|
|
)
|
|
return [_to_watch(row) for row in rows]
|
|
|
|
|
|
async def get_watch(pool: asyncpg.Pool, watch_id: int) -> WatchView | None:
|
|
row = await pool.fetchrow(
|
|
f"SELECT {_WATCH_COLS} FROM watches WHERE id = $1", # noqa: S608
|
|
watch_id,
|
|
)
|
|
return _to_watch(row) if row else None
|
|
|
|
|
|
async def create_watch(
|
|
pool: asyncpg.Pool,
|
|
account_id: int,
|
|
kind: str,
|
|
params: dict[str, Any],
|
|
*,
|
|
enabled: bool = True,
|
|
) -> WatchView:
|
|
row = await pool.fetchrow(
|
|
"INSERT INTO watches (account_id, kind, params, enabled) " # noqa: S608
|
|
f"VALUES ($1, $2, $3::jsonb, $4) RETURNING {_WATCH_COLS}",
|
|
account_id,
|
|
kind,
|
|
json.dumps(params),
|
|
enabled,
|
|
)
|
|
await pool.execute(f"NOTIFY {WATCHES_CHANGED_CHANNEL}")
|
|
return _to_watch(row)
|
|
|
|
|
|
async def update_watch(
|
|
pool: asyncpg.Pool, watch_id: int, params: dict[str, Any], *, enabled: bool
|
|
) -> WatchView | None:
|
|
row = await pool.fetchrow(
|
|
"UPDATE watches SET params = $2::jsonb, enabled = $3, updated_at = now() " # noqa: S608
|
|
f"WHERE id = $1 RETURNING {_WATCH_COLS}",
|
|
watch_id,
|
|
json.dumps(params),
|
|
enabled,
|
|
)
|
|
if row is None:
|
|
return None
|
|
await pool.execute(f"NOTIFY {WATCHES_CHANGED_CHANNEL}")
|
|
return _to_watch(row)
|
|
|
|
|
|
async def delete_watch(pool: asyncpg.Pool, watch_id: int) -> bool:
|
|
result = await pool.execute("DELETE FROM watches WHERE id = $1", watch_id)
|
|
if not result.endswith("1"):
|
|
return False
|
|
await pool.execute(f"NOTIFY {WATCHES_CHANGED_CHANNEL}")
|
|
return True
|
|
|
|
|
|
async def list_alerts(
|
|
pool: asyncpg.Pool, account_id: int, page: Page, *, seen: bool | None = None
|
|
) -> list[AlertView]:
|
|
params: list[object] = [account_id]
|
|
where = "account_id = $1"
|
|
if seen is not None:
|
|
params.append(seen)
|
|
where += f" AND seen = ${len(params)}"
|
|
params.append(page.capped_limit)
|
|
params.append(page.offset)
|
|
rows = await pool.fetch(
|
|
f"SELECT {_ALERT_COLS} FROM alerts WHERE {where} " # noqa: S608
|
|
f"ORDER BY ts DESC LIMIT ${len(params) - 1} OFFSET ${len(params)}",
|
|
*params,
|
|
)
|
|
return [_to_alert(row) for row in rows]
|
|
|
|
|
|
async def insert_alert(
|
|
pool: asyncpg.Pool,
|
|
account_id: int,
|
|
watch_id: int,
|
|
payload: dict[str, Any],
|
|
*,
|
|
dedup_key: str | None = None,
|
|
) -> AlertView | None:
|
|
row = await pool.fetchrow(
|
|
"INSERT INTO alerts (account_id, watch_id, ts, payload, seen, dedup_key) " # noqa: S608
|
|
"VALUES ($1, $2, $3, $4::jsonb, false, $5) "
|
|
"ON CONFLICT (watch_id, dedup_key) WHERE dedup_key IS NOT NULL "
|
|
f"DO NOTHING RETURNING {_ALERT_COLS}",
|
|
account_id,
|
|
watch_id,
|
|
datetime.now(UTC),
|
|
json.dumps(payload),
|
|
dedup_key,
|
|
)
|
|
if row is None:
|
|
return None
|
|
await pool.execute(f"NOTIFY {ALERTS_CHANGED_CHANNEL}")
|
|
return _to_alert(row)
|
|
|
|
|
|
async def mark_alert_seen(pool: asyncpg.Pool, alert_id: int) -> bool:
|
|
result = await pool.execute("UPDATE alerts SET seen = true WHERE id = $1", alert_id)
|
|
return result.endswith("1")
|