from datetime import UTC, datetime from uuid import UUID, uuid4 from fastapi import HTTPException from sqlalchemy import func, select, text from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.ext.asyncio import AsyncSession from .. import metrics from ..agent_auth import AgentCredential from ..logging_config import get_logger from ..models import ( Agent, AgentBlacklist, Check, Event, EventIngestDedup, Incident, NotificationOutbox, ) from ..redaction import redact from ..schemas import CheckResultEvent, HeartbeatRequest from ..settings import get_settings from ..timeutil import to_utc INCIDENT_KEY_MAX = 256 log = get_logger("monlet.ingestion") _SEVERITY_RANK = {"unknown": 1, "warning": 2, "critical": 3} def _severity_for_status(status: str) -> str: if status in ("warning", "critical", "unknown"): return status return "unknown" def _incident_key(agent_id: str, ev: CheckResultEvent) -> str: key = ev.incident_key or f"{agent_id}:{ev.check_id}" if len(key) > INCIDENT_KEY_MAX: raise HTTPException( status_code=400, detail=f"incident_key exceeds {INCIDENT_KEY_MAX} chars", ) return key async def _blacklist_row(session: AsyncSession, cred: AgentCredential) -> AgentBlacklist | None: res = await session.execute( select(AgentBlacklist).where(AgentBlacklist.key_hash == cred.key_hash) ) return res.scalar_one_or_none() async def reject_if_blocked( session: AsyncSession, cred: AgentCredential, agent_id: str, hostname: str, observed: datetime | None = None, ) -> None: row = await _blacklist_row(session, cred) if row is None: return raise HTTPException(status_code=403, detail="agent blocked") async def _ensure_pending_capacity(session: AsyncSession) -> None: max_pending = get_settings().max_pending_agents if max_pending <= 0: return count = ( await session.execute( select(func.count()).select_from(Agent).where(Agent.pending_key_hash.is_not(None)) ) ).scalar_one() if count >= max_pending: raise HTTPException(status_code=429, detail="too many pending agents") async def _reject_hostname_collision(session: AsyncSession, agent_id: str, hostname: str) -> None: row = ( await session.execute( select(Agent).where(Agent.hostname == hostname, Agent.agent_id != agent_id) ) ).scalar_one_or_none() if row is not None: raise HTTPException(status_code=409, detail="hostname already belongs to another agent") def _set_pending( agent: Agent, hb: HeartbeatRequest, cred: AgentCredential, observed: datetime, features: dict, ) -> None: agent.pending_key_hash = cred.key_hash agent.pending_key_fingerprint = cred.key_fingerprint agent.pending_hostname = hb.hostname agent.pending_seen_at = observed agent.pending_features = features agent.pending_labels = hb.labels async def upsert_agent_heartbeat( session: AsyncSession, hb: HeartbeatRequest, cred: AgentCredential ) -> str: observed = to_utc(hb.observed_at) await reject_if_blocked(session, cred, hb.agent_id, hb.hostname, observed) features = hb.features.model_dump() res = await session.execute(select(Agent).where(Agent.accepted_key_hash == cred.key_hash)) accepted = res.scalar_one_or_none() if accepted is not None: await _reject_hostname_collision(session, accepted.agent_id, hb.hostname) accepted.hostname = hb.hostname accepted.status = "alive" accepted.last_seen_at = observed accepted.features = features accepted.labels = hb.labels metrics.heartbeats_total.labels(result="ok").inc() return "accepted" res = await session.execute(select(Agent).where(Agent.pending_key_hash == cred.key_hash)) pending = res.scalar_one_or_none() if pending is not None: await _reject_hostname_collision(session, pending.agent_id, hb.hostname) _set_pending(pending, hb, cred, observed, features) if pending.accepted_key_hash is None: pending.hostname = hb.hostname pending.status = "alive" pending.last_seen_at = observed pending.features = features pending.labels = hb.labels metrics.heartbeats_total.labels(result="pending").inc() return "pending" hostname_row = ( await session.execute(select(Agent).where(Agent.hostname == hb.hostname)) ).scalar_one_or_none() if hostname_row is not None: if hostname_row.pending_key_hash is not None: raise HTTPException( status_code=409, detail="agent already has a different pending key", ) await _ensure_pending_capacity(session) _set_pending(hostname_row, hb, cred, observed, features) metrics.heartbeats_total.labels(result="pending").inc() return "pending" agent_id_row = ( await session.execute(select(Agent).where(Agent.agent_id == hb.agent_id)) ).scalar_one_or_none() if agent_id_row is not None: raise HTTPException(status_code=409, detail="agent_id already exists") await _ensure_pending_capacity(session) session.add( Agent( agent_id=hb.agent_id, hostname=hb.hostname, status="alive", last_seen_at=observed, features=features, labels=hb.labels, pending_key_hash=cred.key_hash, pending_key_fingerprint=cred.key_fingerprint, pending_hostname=hb.hostname, pending_seen_at=observed, pending_features=features, pending_labels=hb.labels, ) ) metrics.heartbeats_total.labels(result="pending").inc() return "pending" async def accepted_agent_for_key(session: AsyncSession, cred: AgentCredential) -> Agent | None: if await _blacklist_row(session, cred) is not None: raise HTTPException(status_code=403, detail="agent blocked") res = await session.execute(select(Agent).where(Agent.accepted_key_hash == cred.key_hash)) return res.scalar_one_or_none() async def _apply_incident( session: AsyncSession, agent_id: str, ev: CheckResultEvent, incident_key: str, ) -> tuple[bool, Incident | None]: """Return (transition, incident). transition True if open->resolved or none->open.""" observed = to_utc(ev.observed_at) if ev.status == "ok": res = await session.execute( select(Incident).where(Incident.incident_key == incident_key, Incident.state == "open") ) inc = res.scalar_one_or_none() if inc is None: return False, None inc.state = "resolved" inc.resolved_at = observed inc.last_event_id = UUID(ev.event_id) metrics.incidents_resolved_total.inc() return True, inc severity = _severity_for_status(ev.status) new_id = uuid4() summary = (redact(ev.output) or "")[:200] if ev.output else None stmt = pg_insert(Incident).values( id=new_id, incident_key=incident_key, agent_id=agent_id, check_id=ev.check_id, state="open", severity=severity, opened_at=observed, last_event_id=UUID(ev.event_id), summary=summary, ) stmt = stmt.on_conflict_do_nothing( index_elements=[Incident.incident_key], index_where=text("state = 'open'"), ).returning(Incident.id) ins = await session.execute(stmt) inserted_id = ins.scalar_one_or_none() if inserted_id is not None: await session.flush() res = await session.execute(select(Incident).where(Incident.id == inserted_id)) metrics.incidents_opened_total.inc() return True, res.scalar_one() res = await session.execute( select(Incident).where(Incident.incident_key == incident_key, Incident.state == "open") ) inc = res.scalar_one() if inc.agent_id != agent_id or inc.check_id != ev.check_id: raise HTTPException( status_code=400, detail="incident_key already bound to another (agent_id, check_id)", ) cur = _SEVERITY_RANK.get(inc.severity, 0) new = _SEVERITY_RANK.get(severity, 0) if new > cur: inc.severity = severity inc.last_event_id = UUID(ev.event_id) return False, inc async def _enqueue_outbox( session: AsyncSession, incident: Incident, event_type: str, payload: dict, ) -> None: notifiers = get_settings().enabled_notifiers if not notifiers: return now = datetime.now(UTC) for name in notifiers: session.add( NotificationOutbox( id=uuid4(), incident_id=incident.id, notifier=name, event_type=event_type, state="pending", attempts=0, next_attempt_at=now, payload=payload, ) ) async def ingest_event( session: AsyncSession, agent_id: str, ev: CheckResultEvent, ) -> bool: """Return True if a new change row was written, False if not (duplicate/unchanged/late). Idempotency is enforced first via the standalone event_ingest_dedup table — a duplicate event_id never touches checks/incidents/outbox/events. The full pipeline (dedup + events + check + incident + outbox) commits as a single transaction by the caller, so a partial failure cannot leave a dedup row without its side-effects. """ observed = to_utc(ev.observed_at) output = redact(ev.output) check_summary = (output or "")[:200] if output else None incident_key = _incident_key(agent_id, ev) event_uuid = UUID(ev.event_id) # Order matters: late check before dedup insert so that repeated late events stay # observable as `late`, not silently collapsed into `deduplicated`. chk_res = await session.execute( select(Check).where(Check.agent_id == agent_id, Check.check_id == ev.check_id) ) cur = chk_res.scalar_one_or_none() is_late = cur is not None and cur.last_observed_at > observed if is_late: metrics.events_total.labels(result="late").inc() return False dedup_stmt = ( pg_insert(EventIngestDedup) .values(event_id=event_uuid, received_at=datetime.now(UTC)) .on_conflict_do_nothing(index_elements=[EventIngestDedup.event_id]) .returning(EventIngestDedup.event_id) ) if (await session.execute(dedup_stmt)).scalar_one_or_none() is None: metrics.events_total.labels(result="deduplicated").inc() return False changed = cur is None or cur.status != ev.status or cur.exit_code != ev.exit_code if changed: stmt = pg_insert(Event).values( event_id=event_uuid, agent_id=agent_id, check_id=ev.check_id, observed_at=observed, status=ev.status, exit_code=ev.exit_code, duration_ms=ev.duration_ms, output=output, output_truncated=ev.output_truncated, notifications_enabled=ev.notifications_enabled, incident_key=incident_key, ) await session.execute(stmt) metrics.events_total.labels(result="accepted").inc() else: metrics.events_total.labels(result="unchanged").inc() # Upsert Check to be safe under concurrent first-write race for the same (agent_id, check_id). # When changed=False keep the existing last_event_id and incident_key — incident_key changes # are only meaningful at transitions and would otherwise desync from any open Incident. check_values = { "agent_id": agent_id, "check_id": ev.check_id, "status": ev.status, "exit_code": ev.exit_code, "last_observed_at": observed, "last_event_id": UUID(ev.event_id), "incident_key": incident_key, "summary": check_summary, } update_set = { "status": ev.status, "exit_code": ev.exit_code, "last_observed_at": observed, "summary": check_summary, } if changed: update_set["last_event_id"] = event_uuid update_set["incident_key"] = incident_key chk_stmt = pg_insert(Check).values(**check_values) chk_stmt = chk_stmt.on_conflict_do_update( index_elements=[Check.agent_id, Check.check_id], set_=update_set ) await session.execute(chk_stmt) if not changed: return False transition, inc = await _apply_incident(session, agent_id, ev, incident_key) if transition and inc is not None and ev.notifications_enabled: event_type = "resolved" if ev.status == "ok" else "firing" await _enqueue_outbox( session, inc, event_type, { "incident_id": str(inc.id), "agent_id": agent_id, "check_id": ev.check_id, "status": ev.status, "severity": inc.severity, "incident_key": incident_key, "observed_at": observed.isoformat(), "opened_at": inc.opened_at.isoformat() if inc.opened_at else None, "resolved_at": inc.resolved_at.isoformat() if inc.resolved_at else None, "summary": inc.summary, "output": output, "event_type": event_type, }, ) return True