monlet/server/monlet_server/models.py

from datetime import datetime
from uuid import UUID

from sqlalchemy import (
    JSON,
    Boolean,
    CheckConstraint,
    DateTime,
    ForeignKey,
    Index,
    Integer,
    PrimaryKeyConstraint,
    String,
    Text,
    func,
    text,
)
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.dialects.postgresql import UUID as PG_UUID
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column


class Base(DeclarativeBase):
    pass


def _jsonb():
    return JSONB().with_variant(JSON(), "sqlite")


def _uuid():
    return PG_UUID(as_uuid=True).with_variant(String(36), "sqlite")


class Agent(Base):
    __tablename__ = "agents"

    agent_id: Mapped[str] = mapped_column(Text, primary_key=True)
    hostname: Mapped[str] = mapped_column(Text, nullable=False)
    status: Mapped[str] = mapped_column(Text, nullable=False, default="alive")
    last_seen_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
    features: Mapped[dict] = mapped_column(
        _jsonb(),
        nullable=False,
        default=lambda: {"push": False, "metrics": False},
        server_default=text("jsonb_build_object('push', false, 'metrics', false)"),
    )
    labels: Mapped[dict] = mapped_column(
        _jsonb(), nullable=False, default=dict, server_default=text("'{}'::jsonb")
    )
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )

    __table_args__ = (
        CheckConstraint("status IN ('alive','stale','dead')", name="agents_status_chk"),
        Index("agents_status_idx", "status"),
        Index("agents_last_seen_idx", "last_seen_at"),
    )


class Check(Base):
    __tablename__ = "checks"

    agent_id: Mapped[str] = mapped_column(
        Text, ForeignKey("agents.agent_id", ondelete="CASCADE"), nullable=False
    )
    check_id: Mapped[str] = mapped_column(Text, nullable=False)
    status: Mapped[str] = mapped_column(Text, nullable=False)
    exit_code: Mapped[int] = mapped_column(Integer, nullable=False)
    last_observed_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
    last_event_id: Mapped[UUID] = mapped_column(_uuid(), nullable=False)
    incident_key: Mapped[str] = mapped_column(Text, nullable=False)

    __table_args__ = (
        PrimaryKeyConstraint("agent_id", "check_id", name="checks_pkey"),
        CheckConstraint(
            "status IN ('ok','warning','critical','unknown')", name="checks_status_chk"
        ),
        Index("checks_status_idx", "status"),
    )


class Event(Base):
    __tablename__ = "events"

    event_id: Mapped[UUID] = mapped_column(_uuid(), primary_key=True)
    agent_id: Mapped[str] = mapped_column(Text, nullable=False)
    check_id: Mapped[str] = mapped_column(Text, nullable=False)
    observed_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
    received_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    status: Mapped[str] = mapped_column(Text, nullable=False)
    exit_code: Mapped[int] = mapped_column(Integer, nullable=False)
    duration_ms: Mapped[int] = mapped_column(Integer, nullable=False)
    output: Mapped[str | None] = mapped_column(Text, nullable=True)
    output_truncated: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
    notifications_enabled: Mapped[bool] = mapped_column(
        Boolean, nullable=False, default=True, server_default=text("TRUE")
    )
    incident_key: Mapped[str] = mapped_column(Text, nullable=False)

    __table_args__ = (
        CheckConstraint(
            "status IN ('ok','warning','critical','unknown')", name="events_status_chk"
        ),
        CheckConstraint("duration_ms >= 0", name="events_duration_chk"),
        Index(
            "events_agent_check_observed_idx",
            "agent_id",
            "check_id",
            text("observed_at DESC"),
        ),
        Index("events_observed_at_idx", text("observed_at DESC")),
    )


class Incident(Base):
    __tablename__ = "incidents"

    id: Mapped[UUID] = mapped_column(_uuid(), primary_key=True)
    incident_key: Mapped[str] = mapped_column(Text, nullable=False)
    agent_id: Mapped[str] = mapped_column(Text, nullable=False)
    check_id: Mapped[str] = mapped_column(Text, nullable=False)
    state: Mapped[str] = mapped_column(Text, nullable=False)
    severity: Mapped[str] = mapped_column(Text, nullable=False)
    opened_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
    resolved_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    summary: Mapped[str | None] = mapped_column(Text, nullable=True)
    last_event_id: Mapped[UUID] = mapped_column(_uuid(), nullable=False)

    __table_args__ = (
        CheckConstraint("state IN ('open','resolved')", name="incidents_state_chk"),
        CheckConstraint(
            "severity IN ('warning','critical','unknown')", name="incidents_severity_chk"
        ),
        Index(
            "incidents_open_key_idx",
            "incident_key",
            unique=True,
            postgresql_where=text("state = 'open'"),
        ),
        Index("incidents_state_opened_idx", "state", text("opened_at DESC")),
    )


class NotificationOutbox(Base):
    __tablename__ = "notification_outbox"

    id: Mapped[UUID] = mapped_column(_uuid(), primary_key=True)
    incident_id: Mapped[UUID] = mapped_column(
        _uuid(), ForeignKey("incidents.id", ondelete="CASCADE"), nullable=False
    )
    notifier: Mapped[str] = mapped_column(Text, nullable=False)
    event_type: Mapped[str] = mapped_column(Text, nullable=False)
    state: Mapped[str] = mapped_column(Text, nullable=False)
    attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    next_attempt_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    last_error: Mapped[str | None] = mapped_column(Text, nullable=True)
    payload: Mapped[dict] = mapped_column(_jsonb(), nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )

    __table_args__ = (
        CheckConstraint("event_type IN ('firing','resolved')", name="outbox_event_type_chk"),
        CheckConstraint(
            "state IN ('pending','sending','sent','retry','failed','discarded')",
            name="outbox_state_chk",
        ),
        Index(
            "outbox_pending_idx",
            "state",
            "next_attempt_at",
            postgresql_where=text("state IN ('pending','retry')"),
        ),
    )