AWS-Panel/backend/modules/jobs/models.py

from __future__ import annotations

from datetime import datetime
from enum import Enum
from typing import Optional

from sqlalchemy import DateTime, Enum as SAEnum, ForeignKey, Index, JSON, String, UniqueConstraint, text
from sqlalchemy.dialects.mysql import BIGINT, INTEGER, TINYINT
from sqlalchemy.orm import Mapped, mapped_column, relationship

from backend.db.base import Base


class JobType(str, Enum):
    SYNC_INSTANCES = "SYNC_INSTANCES"
    START_INSTANCES = "START_INSTANCES"
    STOP_INSTANCES = "STOP_INSTANCES"
    REBOOT_INSTANCES = "REBOOT_INSTANCES"
    TERMINATE_INSTANCES = "TERMINATE_INSTANCES"
    CREATE_INSTANCES = "CREATE_INSTANCES"


class JobStatus(str, Enum):
    PENDING = "PENDING"
    RUNNING = "RUNNING"
    SUCCESS = "SUCCESS"
    FAILED = "FAILED"


class JobItemResourceType(str, Enum):
    INSTANCE = "INSTANCE"
    OTHER = "OTHER"


class JobItemAction(str, Enum):
    CREATE = "CREATE"
    START = "START"
    STOP = "STOP"
    REBOOT = "REBOOT"
    TERMINATE = "TERMINATE"
    SYNC = "SYNC"


class JobItemStatus(str, Enum):
    PENDING = "PENDING"
    RUNNING = "RUNNING"
    SUCCESS = "SUCCESS"
    FAILED = "FAILED"
    SKIPPED = "SKIPPED"


class Job(Base):
    __tablename__ = "jobs"
    __table_args__ = (
        UniqueConstraint("job_uuid", name="uniq_job_uuid"),
        Index("idx_jobs_type", "job_type"),
        Index("idx_jobs_status", "status"),
        Index("idx_jobs_created_at", "created_at"),
    )

    id: Mapped[int] = mapped_column(BIGINT(unsigned=True), primary_key=True, autoincrement=True)
    job_uuid: Mapped[str] = mapped_column(String(32), nullable=False)
    job_type: Mapped[JobType] = mapped_column(SAEnum(JobType), nullable=False)
    status: Mapped[JobStatus] = mapped_column(SAEnum(JobStatus), nullable=False, server_default=text("'PENDING'"))
    progress: Mapped[int] = mapped_column(TINYINT(unsigned=True), nullable=False, server_default=text("0"))
    total_count: Mapped[Optional[int]] = mapped_column(INTEGER(unsigned=True), server_default=text("0"))
    success_count: Mapped[Optional[int]] = mapped_column(INTEGER(unsigned=True), server_default=text("0"))
    fail_count: Mapped[Optional[int]] = mapped_column(INTEGER(unsigned=True), server_default=text("0"))
    skipped_count: Mapped[Optional[int]] = mapped_column(INTEGER(unsigned=True), server_default=text("0"))
    payload: Mapped[Optional[dict]] = mapped_column(JSON)
    error_message: Mapped[Optional[str]] = mapped_column(String(512))
    created_by_user_id: Mapped[Optional[int]] = mapped_column(
        ForeignKey("users.id", ondelete="SET NULL", onupdate="CASCADE")
    )
    created_for_customer: Mapped[Optional[int]] = mapped_column(
        ForeignKey("customers.id", ondelete="SET NULL", onupdate="CASCADE")
    )
    created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP"), nullable=False)
    started_at: Mapped[Optional[datetime]] = mapped_column(DateTime)
    finished_at: Mapped[Optional[datetime]] = mapped_column(DateTime)
    updated_at: Mapped[datetime] = mapped_column(
        DateTime, server_default=text("CURRENT_TIMESTAMP"), onupdate=text("CURRENT_TIMESTAMP"), nullable=False
    )

    items: Mapped[list["JobItem"]] = relationship("JobItem", back_populates="job")


class JobItem(Base):
    __tablename__ = "job_items"
    __table_args__ = (
        Index("idx_job_items_job", "job_id"),
        Index("idx_job_items_instance", "resource_id"),
        Index("idx_job_items_status", "status"),
    )

    id: Mapped[int] = mapped_column(BIGINT(unsigned=True), primary_key=True, autoincrement=True)
    job_id: Mapped[int] = mapped_column(ForeignKey("jobs.id", ondelete="CASCADE", onupdate="CASCADE"), nullable=False)
    resource_type: Mapped[JobItemResourceType] = mapped_column(
        SAEnum(JobItemResourceType), nullable=False, server_default=text("'INSTANCE'")
    )
    resource_id: Mapped[Optional[int]] = mapped_column(
        ForeignKey("instances.id", ondelete="SET NULL", onupdate="CASCADE")
    )
    account_id: Mapped[Optional[str]] = mapped_column(String(32))
    region: Mapped[Optional[str]] = mapped_column(String(32))
    instance_id: Mapped[Optional[str]] = mapped_column(String(32))
    action: Mapped[JobItemAction] = mapped_column(SAEnum(JobItemAction), nullable=False)
    status: Mapped[JobItemStatus] = mapped_column(
        SAEnum(JobItemStatus), nullable=False, server_default=text("'PENDING'")
    )
    error_message: Mapped[Optional[str]] = mapped_column(String(512))
    extra: Mapped[Optional[dict]] = mapped_column(JSON)
    created_at: Mapped[datetime] = mapped_column(DateTime, server_default=text("CURRENT_TIMESTAMP"), nullable=False)
    updated_at: Mapped[datetime] = mapped_column(
        DateTime, server_default=text("CURRENT_TIMESTAMP"), onupdate=text("CURRENT_TIMESTAMP"), nullable=False
    )

    job: Mapped["Job"] = relationship("Job", back_populates="items")
    instance: Mapped[Optional["Instance"]] = relationship("Instance", back_populates="job_items")