argos/argos_monitoring/server/models.py

"""Database models"""

from datetime import datetime, timedelta
from typing import List, Literal

from sqlalchemy import (
    JSON,
    Enum,
    ForeignKey,
)
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship

from argos_monitoring.checks import BaseCheck, get_registered_check
from argos_monitoring.schemas import WebsiteCheck


class Base(DeclarativeBase):
    type_annotation_map = {List[WebsiteCheck]: JSON, dict: JSON}


class Task(Base):
    """
    There is one task per check.

    It contains all information needed to run the jobs on the agents.
    Agents will return information in the result table.
    """

    __tablename__ = "tasks"
    id: Mapped[int] = mapped_column(primary_key=True)

    # Info needed to run the task
    url: Mapped[str] = mapped_column()
    domain: Mapped[str] = mapped_column()
    check: Mapped[str] = mapped_column()
    expected: Mapped[str] = mapped_column()
    frequency: Mapped[int] = mapped_column()

    # Orchestration-related
    selected_by: Mapped[str] = mapped_column(nullable=True)
    selected_at: Mapped[datetime] = mapped_column(nullable=True)
    completed_at: Mapped[datetime] = mapped_column(nullable=True)
    next_run: Mapped[datetime] = mapped_column(nullable=True)

    severity: Mapped[Literal["ok", "warning", "critical", "unknown"]] = mapped_column(
        Enum("ok", "warning", "critical", "unknown", name="severity"),
        insert_default="unknown",
    )
    last_severity_update: Mapped[datetime] = mapped_column(nullable=True)

    results: Mapped[List["Result"]] = relationship(
        back_populates="task",
        cascade="all, delete",
        passive_deletes=True,
    )

    def __str__(self):
        return f"DB Task {self.url} - {self.check} - {self.expected}"

    def get_check(self) -> BaseCheck:
        """Returns a check instance for this specific task"""
        return get_registered_check(self.check)

    def set_times_severity_and_deselect(self, severity, submitted_at):
        """Removes the lock on task, set its severity and set the time for the next run"""
        self.severity = severity
        self.last_severity_update = submitted_at
        self.selected_by = None
        self.selected_at = None

        now = datetime.now()
        self.completed_at = now
        self.next_run = now + timedelta(minutes=self.frequency)

    @property
    def last_result(self):
        """Get last result of the task"""
        if not self.results:
            return None
        return max(self.results, key=lambda r: r.id)

    @property
    def status(self):
        """Get status of the task, i.e. the status of its last result"""
        if not self.last_result:
            return None
        return self.last_result.status


class Result(Base):
    """There are multiple results per task.

    The results store information returned by the agents.

    You can read `status` as  "Was the agent able to do the check?"
    while the `severity` depends on the return value of the check.
    """

    __tablename__ = "results"
    id: Mapped[int] = mapped_column(primary_key=True)
    task_id: Mapped[int] = mapped_column(ForeignKey("tasks.id", ondelete="CASCADE"))
    task: Mapped["Task"] = relationship(back_populates="results")
    agent_id: Mapped[str] = mapped_column(nullable=True)

    submitted_at: Mapped[datetime] = mapped_column()
    # XXX change "on-check" to something better.
    status: Mapped[Literal["success", "failure", "error", "on-check"]] = mapped_column(
        Enum("success", "failure", "error", "on-check", name="status")
    )
    severity: Mapped[Literal["ok", "warning", "critical", "unknown"]] = mapped_column(
        Enum("ok", "warning", "critical", "unknown", name="severity")
    )
    context: Mapped[dict] = mapped_column()

    def set_status(self, status, severity):
        self.severity = severity
        self.status = status

    def __str__(self):
        return f"DB Result {self.id} - {self.status} - {self.context}"


class ConfigCache(Base):
    """Database model containing information on the current state
    of the configuration.

    This is used to determine if tasks are to be updated.

    These settings are cached:
    - general_frequency: the content of general.frequency setting, in minutes
                         ex: 5
    - websites_hash:     the hash (sha256sum) of websites setting, to allow a quick
                         comparison without looping through all websites.
                         ex: 8b886e7db7b553fe99f6d5437f31745987e243c77b2109b84cf9a7f8bf7d75b1
    """

    __tablename__ = "config_cache"
    name: Mapped[str] = mapped_column(primary_key=True)
    val: Mapped[str] = mapped_column()
    updated_at: Mapped[datetime] = mapped_column()


class User(Base):
    """Database model for user authentication"""

    __tablename__ = "users"
    username: Mapped[str] = mapped_column(primary_key=True)
    password: Mapped[str] = mapped_column()
    disabled: Mapped[bool] = mapped_column()
    created_at: Mapped[datetime] = mapped_column(default=datetime.now())
    updated_at: Mapped[datetime] = mapped_column(nullable=True)
    last_login_at: Mapped[datetime] = mapped_column(nullable=True)

    def update_last_login_at(self):
        self.last_login_at = datetime.now()