from fastapi import APIRouter, Depends, HTTPException from uuid import UUID import json # ← REQUIRED for json.loads from backend.app.dependencies import validate_token from backend.app.tasks.schema import ( CreateTaskRequest, UpdateTaskStatusRequest, CreateSubTaskRequest, UpdateSubTaskStatusRequest, FinalizeSubTaskRequest, TaskSqsTriggerRequest, ) # Correct location of interfaces from backend.app.db.functions.tasks.Tasks import TasksInterface, SubTaskInterface from backend.app.db.connection import get_db_session from backend.app.db.models.tasks import Task, SubTask from sqlmodel import select router = APIRouter( prefix="/tasks", tags=["tasks"], dependencies=[Depends(validate_token)], ) # ============================================================ # Create Task # ============================================================ @router.post("/", summary="Create a new task and its first subtask") async def create_task(req: CreateTaskRequest): tasks = TasksInterface() task_id, subtask_id = tasks.create_task( task_source=req.task_source, service=req.service, inputs=req.inputs, ) return {"task_id": task_id, "subtask_id": subtask_id} # ============================================================ # Get Task + Subtasks # ============================================================ @router.get("/{task_id}", summary="Get a task and its subtasks") async def get_task(task_id: UUID): with get_db_session() as session: task = session.get(Task, task_id) if not task: raise HTTPException(status_code=404, detail="Task not found") subtasks = session.exec(select(SubTask).where(SubTask.taskId == task_id)).all() formatted = [] for st in subtasks: formatted.append( { **st.dict(), "inputs": json.loads(st.inputs) if st.inputs else None, "outputs": json.loads(st.outputs) if st.outputs else None, "cloud_logs_url": st.cloudLogsURL, } ) return { "task": task, "subtasks": formatted, } # ============================================================ # Update Task Status # ============================================================ @router.put("/{task_id}/status", summary="Update a task's status") async def update_task_status(task_id: UUID, req: UpdateTaskStatusRequest): tasks = TasksInterface() try: updated = tasks.update_task_status(task_id, req.status) return {"task_id": updated.id, "status": updated.status} except ValueError as e: raise HTTPException(status_code=404, detail=str(e)) # ============================================================ # Create Additional Subtask # ============================================================ @router.post("/{task_id}/subtasks", summary="Create a new subtask under a task") async def create_subtask(task_id: UUID, req: CreateSubTaskRequest): subtasks = SubTaskInterface() try: st = subtasks.create_subtask(task_id, req.inputs) return {"subtask_id": st.id, "task_id": task_id, "status": st.status} except ValueError as e: raise HTTPException(status_code=404, detail=str(e)) # ============================================================ # Update Subtask Status # ============================================================ @router.put("/subtask/{subtask_id}/status", summary="Update a subtask's status") async def update_subtask_status(subtask_id: UUID, req: UpdateSubTaskStatusRequest): subtasks = SubTaskInterface() try: st = subtasks.update_subtask_status(subtask_id, req.status) return {"subtask_id": st.id, "status": st.status} except ValueError as e: raise HTTPException(status_code=404, detail=str(e)) # === # Sub task is complete @router.post( "/subtask/{subtask_id}/finalize", summary="Finalize a subtask with status, outputs, logs", ) async def finalize_subtask(subtask_id: UUID, req: FinalizeSubTaskRequest): subtasks = SubTaskInterface() try: st = subtasks.finalize_subtask( subtask_id=subtask_id, status=req.status, outputs=req.outputs, cloud_logs_url=req.cloud_logs_url, ) return { "subtask_id": st.id, "status": st.status, "outputs": req.outputs, "cloud_logs_url": req.cloud_logs_url, } except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) # for testing: import boto3 import json from backend.app.tasks.schema import TaskSqsTriggerRequest from backend.app.db.functions.tasks.Tasks import TasksInterface, SubTaskInterface from backend.app.config import get_settings @router.post( "/trigger", summary="Create task + subtask and publish to SQS", status_code=202 ) async def trigger_task(req: TaskSqsTriggerRequest): """ Creates a Task + SubTask, then pushes the SubTask into SQS so a Lambda can process it. If inputs are empty, automatically replaced with {}. """ settings = get_settings() sqs = boto3.client("sqs", settings.AWS_DEFAULT_REGION) tasks = TasksInterface() # ---- Normalize empty inputs ---- inputs = req.inputs or {} # ensures {} even if null # ---- 1. Create Task + SubTask ---- task_id, subtask_id = tasks.create_task( task_source=req.task_source, service=req.service, inputs=inputs, ) # ---- 2. Prepare SQS payload ---- sqs_payload = { "subtask_id": str(subtask_id), "params": inputs, } try: response = sqs.send_message( QueueUrl=f"https://sqs.{settings.AWS_REGION}.amazonaws.com/" f"{settings.AWS_ACCOUNT_ID}/lambda-example-queue", MessageBody=json.dumps(sqs_payload), ) except Exception as e: raise HTTPException(status_code=500, detail=f"SQS error: {e}") return { "message": "Task triggered", "task_id": task_id, "subtask_id": subtask_id, "sqs_message_id": response.get("MessageId"), "inputs_sent": inputs, }