argos/argos/server/models.py

72 lines
2.1 KiB
Python

from datetime import datetime
from typing import List, Literal
from sqlalchemy import (
JSON,
Boolean,
Column,
DateTime,
Enum,
ForeignKey,
Integer,
String,
)
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
from sqlalchemy_utils import ChoiceType
from argos.checks import get_registered_check
from argos.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 workers.
Workers 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()
# Orchestration-related
selected_by: Mapped[str] = mapped_column(nullable=True)
selected_at: Mapped[datetime] = mapped_column(nullable=True)
results: Mapped[List["Result"]] = relationship(back_populates="task")
def __str__(self):
return f"DB Task {self.url} - {self.check} - {self.expected}"
def get_check(self):
"""Returns a check instance for this specific task"""
return get_registered_check(self.check)
class Result(Base):
__tablename__ = "results"
id: Mapped[int] = mapped_column(primary_key=True)
task_id: Mapped[int] = mapped_column(ForeignKey("tasks.id"))
task: Mapped["Task"] = relationship(back_populates="results")
submitted_at: Mapped[datetime] = mapped_column()
status: Mapped[Literal["success", "failure", "error", "on-check"]] = mapped_column(
Enum("success", "failure", "error", "on-check")
)
severity: Mapped[Literal["ok", "warning", "critical"]] = mapped_column(
Enum("ok", "warning", "critical")
)
context: Mapped[dict] = mapped_column()
def __str__(self):
return f"DB Result {self.id} - {self.status} - {self.context}"