diff --git a/backend/app/db/functions/tasks/Tasks.py b/backend/app/db/functions/tasks/Tasks.py index 0f987f3b..7ba3dd35 100644 --- a/backend/app/db/functions/tasks/Tasks.py +++ b/backend/app/db/functions/tasks/Tasks.py @@ -11,7 +11,7 @@ from sqlmodel import Session, select from backend.app.db.connection import get_db_session # ---- Models ---- -from backend.app.db.models.tasks import Task, SubTask +from backend.app.db.models.tasks import SourceEnum, Task, SubTask # ============================================================ @@ -268,6 +268,8 @@ class TasksInterface: service: Optional[str] = None, inputs: Optional[Dict[str, Any]] = None, task_only: bool = False, + source: Optional[SourceEnum] = None, + source_id: Optional[str] = None, ): """ Create a new Task record, and an initial SubTask in waiting state. Can also be used to create just @@ -286,6 +288,8 @@ class TasksInterface: status="waiting", job_started=now, job_completed=None, + source=source, + source_id=source_id, ) session.add(task) diff --git a/backend/app/db/models/tasks.py b/backend/app/db/models/tasks.py index cfe18d83..e97a939f 100644 --- a/backend/app/db/models/tasks.py +++ b/backend/app/db/models/tasks.py @@ -1,14 +1,24 @@ +import enum from typing import Optional from datetime import datetime from uuid import UUID, uuid4 +from sqlalchemy import Column, Enum from sqlmodel import SQLModel, Field, Relationship +class SourceEnum(enum.Enum): # TODO: move to domain? + PORTFOLIO = "portfolio_id" + + class Task(SQLModel, table=True): __tablename__ = "tasks" - id: UUID = Field(default_factory=uuid4, primary_key=True, index=True, ) + id: UUID = Field( + default_factory=uuid4, + primary_key=True, + index=True, + ) task_source: str job_started: Optional[datetime] = None job_completed: Optional[datetime] = None @@ -16,13 +26,32 @@ class Task(SQLModel, table=True): service: Optional[str] = None updated_at: datetime = Field(default_factory=datetime.utcnow) + # source: Mapped[Optional[SourceEnum]] = mapped_column(Enum(SourceEnum)) <- SQLAlchemy not SQLModel + + source: Optional[SourceEnum] = Field( + default=None, + sa_column=Column( + Enum( + SourceEnum, + name="source", + values_callable=lambda e: [m.value for m in e], + ), + nullable=True, + ), + ) + source_id: Optional[str] = None + sub_tasks: list["SubTask"] = Relationship(back_populates="task") class SubTask(SQLModel, table=True): __tablename__ = "sub_task" - id: UUID = Field(default_factory=uuid4, primary_key=True, index=True, ) + id: UUID = Field( + default_factory=uuid4, + primary_key=True, + index=True, + ) task_id: UUID = Field(foreign_key="tasks.id") job_started: Optional[datetime] = None diff --git a/backend/app/plan/router.py b/backend/app/plan/router.py index cc9fb579..a90c391a 100644 --- a/backend/app/plan/router.py +++ b/backend/app/plan/router.py @@ -9,6 +9,7 @@ import asyncio from datetime import datetime from fastapi import APIRouter, Depends +from backend.app.db.models.tasks import SourceEnum from backend.app.dependencies import validate_token from backend.app.plan.schemas import PlanTriggerRequest from backend.app.config import get_settings @@ -68,6 +69,8 @@ async def trigger_categorisation( service="plan_categorisation", inputs=payload.model_dump(), task_only=True, + source=SourceEnum.PORTFOLIO, + source_id=str(payload.portfolio_id), ) # Dispatch requests to lambdas