Merge branch 'main' into deploy-backend-with-terraform

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)

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

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