state-hub/api/routers/tasks.py

import uuid

from fastapi import APIRouter, Depends, HTTPException, Query, status
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from api.database import get_session
from api.models.task import Task, TaskStatus
from api.schemas.task import TaskCreate, TaskRead, TaskUpdate

router = APIRouter(prefix="/tasks", tags=["tasks"])


@router.get("/", response_model=list[TaskRead])
async def list_tasks(
    workstream_id: uuid.UUID | None = None,
    status: TaskStatus | None = None,
    assignee: str | None = None,
    needs_human: bool | None = Query(None),
    session: AsyncSession = Depends(get_session),
) -> list[Task]:
    q = select(Task)
    if workstream_id:
        q = q.where(Task.workstream_id == workstream_id)
    if status:
        q = q.where(Task.status == status)
    if assignee:
        q = q.where(Task.assignee == assignee)
    if needs_human is not None:
        q = q.where(Task.needs_human == needs_human)
    q = q.order_by(Task.created_at)
    result = await session.execute(q)
    return list(result.scalars().all())


@router.post("/", response_model=TaskRead, status_code=status.HTTP_201_CREATED)
async def create_task(
    body: TaskCreate,
    session: AsyncSession = Depends(get_session),
) -> Task:
    task = Task(**body.model_dump())
    session.add(task)
    await session.commit()
    await session.refresh(task)
    return task


@router.get("/{task_id}", response_model=TaskRead)
async def get_task(
    task_id: uuid.UUID,
    session: AsyncSession = Depends(get_session),
) -> Task:
    task = await session.get(Task, task_id)
    if task is None:
        raise HTTPException(status_code=404, detail="Task not found")
    return task


@router.patch("/{task_id}", response_model=TaskRead)
async def update_task(
    task_id: uuid.UUID,
    body: TaskUpdate,
    session: AsyncSession = Depends(get_session),
) -> Task:
    task = await session.get(Task, task_id)
    if task is None:
        raise HTTPException(status_code=404, detail="Task not found")
    for field, value in body.model_dump(exclude_unset=True).items():
        setattr(task, field, value)
    await session.commit()
    await session.refresh(task)
    return task


@router.delete("/{task_id}", response_model=TaskRead)
async def cancel_task(
    task_id: uuid.UUID,
    session: AsyncSession = Depends(get_session),
) -> Task:
    task = await session.get(Task, task_id)
    if task is None:
        raise HTTPException(status_code=404, detail="Task not found")
    task.status = TaskStatus.cancelled
    await session.commit()
    await session.refresh(task)
    return task