86 lines
2.2 KiB
Python
86 lines
2.2 KiB
Python
"""Task model for orchestration jobs."""
|
|
|
|
import uuid
|
|
from enum import Enum
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from sqlalchemy import ForeignKey, JSON, String
|
|
from sqlalchemy.dialects.postgresql import JSONB
|
|
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
|
|
|
# Use JSONB on PostgreSQL, JSON on other databases (SQLite for tests)
|
|
JSONType = JSON().with_variant(JSONB, "postgresql")
|
|
|
|
from app.models.base import Base, TimestampMixin, UUIDMixin
|
|
|
|
if TYPE_CHECKING:
|
|
from app.models.agent import Agent
|
|
from app.models.event import Event
|
|
from app.models.tenant import Tenant
|
|
|
|
|
|
class TaskStatus(str, Enum):
|
|
"""Task execution status."""
|
|
|
|
PENDING = "pending"
|
|
RUNNING = "running"
|
|
COMPLETED = "completed"
|
|
FAILED = "failed"
|
|
|
|
|
|
class Task(UUIDMixin, TimestampMixin, Base):
|
|
"""
|
|
Task model representing an orchestration job.
|
|
|
|
Tasks are assigned to agents and track execution state.
|
|
Payload and result use JSONB for flexible, queryable storage.
|
|
"""
|
|
|
|
__tablename__ = "tasks"
|
|
|
|
tenant_id: Mapped[uuid.UUID] = mapped_column(
|
|
ForeignKey("tenants.id", ondelete="CASCADE"),
|
|
nullable=False,
|
|
index=True,
|
|
)
|
|
agent_id: Mapped[uuid.UUID | None] = mapped_column(
|
|
ForeignKey("agents.id", ondelete="SET NULL"),
|
|
nullable=True,
|
|
index=True,
|
|
)
|
|
type: Mapped[str] = mapped_column(
|
|
String(100),
|
|
nullable=False,
|
|
index=True,
|
|
)
|
|
payload: Mapped[dict[str, Any]] = mapped_column(
|
|
JSONType,
|
|
nullable=False,
|
|
default=dict,
|
|
)
|
|
status: Mapped[str] = mapped_column(
|
|
String(50),
|
|
default=TaskStatus.PENDING.value,
|
|
nullable=False,
|
|
index=True,
|
|
)
|
|
result: Mapped[dict[str, Any] | None] = mapped_column(
|
|
JSONType,
|
|
nullable=True,
|
|
)
|
|
|
|
# Relationships
|
|
tenant: Mapped["Tenant"] = relationship(
|
|
back_populates="tasks",
|
|
)
|
|
agent: Mapped["Agent | None"] = relationship(
|
|
back_populates="tasks",
|
|
)
|
|
events: Mapped[list["Event"]] = relationship(
|
|
back_populates="task",
|
|
lazy="selectin",
|
|
)
|
|
|
|
def __repr__(self) -> str:
|
|
return f"<Task(id={self.id}, type={self.type}, status={self.status})>"
|