Repository: PrefectHQ/prefect

Issue:
Port `ui/task_runs/dashboard/counts` endpoint from Cloud to OSS
### First check

- [X] I am a contributor to the Prefect codebase

### Description

In cloud we have a `ui/task_runs/dashboard/counts` endpoint that returns counts of the number of completed and the number of failed task runs that match a filter over time. We should port that to open source so that the upcoming workspace dashboard work can take advantage of it.


### Impact

_No response_

### Additional context

_No response_


Discussion:

Resolving patch:
diff --git a/src/prefect/server/api/server.py b/src/prefect/server/api/server.py
--- a/src/prefect/server/api/server.py
+++ b/src/prefect/server/api/server.py
@@ -83,6 +83,7 @@
     api.collections.router,
     api.variables.router,
     api.ui.flow_runs.router,
+    api.ui.task_runs.router,
     api.admin.router,
     api.root.router,
 )
diff --git a/src/prefect/server/api/ui/__init__.py b/src/prefect/server/api/ui/__init__.py
--- a/src/prefect/server/api/ui/__init__.py
+++ b/src/prefect/server/api/ui/__init__.py
@@ -1,2 +1,3 @@
 """Routes primarily for use by the UI"""
 import prefect.server.api.ui.flow_runs
+import prefect.server.api.ui.task_runs
diff --git a/src/prefect/server/api/ui/task_runs.py b/src/prefect/server/api/ui/task_runs.py
new file mode 100644
--- /dev/null
+++ b/src/prefect/server/api/ui/task_runs.py
@@ -0,0 +1,183 @@
+import sys
+from datetime import datetime, timezone
+from typing import List, Optional, cast
+
+import pendulum
+import sqlalchemy as sa
+from fastapi import Depends, HTTPException, status
+from pydantic import Field
+
+import prefect.server.schemas as schemas
+from prefect.logging import get_logger
+from prefect.server import models
+from prefect.server.database.dependencies import provide_database_interface
+from prefect.server.database.interface import PrefectDBInterface
+from prefect.server.utilities.schemas import DateTimeTZ, PrefectBaseModel
+from prefect.server.utilities.server import PrefectRouter
+
+logger = get_logger("orion.api.ui.task_runs")
+
+router = PrefectRouter(prefix="/ui/task_runs", tags=["Task Runs", "UI"])
+
+FAILED_STATES = [schemas.states.StateType.CRASHED, schemas.states.StateType.FAILED]
+
+
+class TaskRunCount(PrefectBaseModel):
+    completed: int = Field(
+        default=..., description="The number of completed task runs."
+    )
+    failed: int = Field(default=..., description="The number of failed task runs.")
+
+
+def _postgres_bucket_expression(
+    db: PrefectDBInterface, delta: pendulum.Duration, start_datetime: datetime
+):
+    # asyncpg under Python 3.7 doesn't support timezone-aware datetimes for the EXTRACT
+    # function, so we will send it as a naive datetime in UTC
+    if sys.version_info < (3, 8):
+        start_datetime = start_datetime.astimezone(timezone.utc).replace(tzinfo=None)
+
+    return sa.func.floor(
+        (
+            sa.func.extract("epoch", db.TaskRun.start_time)
+            - sa.func.extract("epoch", start_datetime)
+        )
+        / delta.total_seconds()
+    ).label("bucket")
+
+
+def _sqlite_bucket_expression(
+    db: PrefectDBInterface, delta: pendulum.Duration, start_datetime: datetime
+):
+    return sa.func.floor(
+        (
+            (
+                sa.func.strftime("%s", db.TaskRun.start_time)
+                - sa.func.strftime("%s", start_datetime)
+            )
+            / delta.total_seconds()
+        )
+    ).label("bucket")
+
+
+@router.post("/dashboard/counts")
+async def read_dashboard_task_run_counts(
+    task_runs: schemas.filters.TaskRunFilter,
+    flows: Optional[schemas.filters.FlowFilter] = None,
+    flow_runs: Optional[schemas.filters.FlowRunFilter] = None,
+    deployments: Optional[schemas.filters.DeploymentFilter] = None,
+    work_pools: Optional[schemas.filters.WorkPoolFilter] = None,
+    work_queues: Optional[schemas.filters.WorkQueueFilter] = None,
+    db: PrefectDBInterface = Depends(provide_database_interface),
+) -> List[TaskRunCount]:
+    if task_runs.start_time is None or task_runs.start_time.after_ is None:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail="task_runs.start_time.after_ is required",
+        )
+
+    # We only care about task runs that are in a terminal state, all others
+    # should be ignored.
+    task_runs.state = schemas.filters.TaskRunFilterState(
+        type=schemas.filters.TaskRunFilterStateType(
+            any_=list(schemas.states.TERMINAL_STATES)
+        )
+    )
+
+    bucket_count = 20
+    start_time = task_runs.start_time.after_.start_of("minute")
+    end_time = cast(
+        DateTimeTZ,
+        (
+            task_runs.start_time.before_.end_of("minute")
+            if task_runs.start_time.before_
+            else pendulum.now("UTC").end_of("minute")
+        ),
+    )
+    window = end_time - start_time
+    delta = window.as_timedelta() / bucket_count
+
+    async with db.session_context(begin_transaction=False) as session:
+        # Gather the raw counts. The counts are divided into buckets of time
+        # and each bucket contains the number of successful and failed task
+        # runs.
+        # SQLAlchemy doesn't play nicely with our DateTimeTZ type so we convert it
+        # to a datetime object.
+        start_datetime = datetime(
+            start_time.year,
+            start_time.month,
+            start_time.day,
+            start_time.hour,
+            start_time.minute,
+            start_time.second,
+            start_time.microsecond,
+            start_time.timezone,
+        )
+        bucket_expression = (
+            _sqlite_bucket_expression(db, delta, start_datetime)
+            if db.dialect.name == "sqlite"
+            else _postgres_bucket_expression(db, delta, start_datetime)
+        )
+
+        raw_counts = (
+            (
+                await models.task_runs._apply_task_run_filters(
+                    sa.select(
+                        bucket_expression,
+                        sa.func.min(db.TaskRun.end_time).label("oldest"),
+                        sa.func.sum(
+                            sa.case(
+                                (
+                                    db.TaskRun.state_type.in_(FAILED_STATES),
+                                    1,
+                                ),
+                                else_=0,
+                            )
+                        ).label("failed_count"),
+                        sa.func.sum(
+                            sa.case(
+                                (
+                                    db.TaskRun.state_type.notin_(FAILED_STATES),
+                                    1,
+                                ),
+                                else_=0,
+                            )
+                        ).label("successful_count"),
+                    ),
+                    flow_filter=flows,
+                    flow_run_filter=flow_runs,
+                    task_run_filter=task_runs,
+                    deployment_filter=deployments,
+                    work_pool_filter=work_pools,
+                    work_queue_filter=work_queues,
+                )
+            )
+            .group_by("bucket", db.TaskRun.start_time)
+            .subquery()
+        )
+
+        # Aggregate the raw counts by bucket
+        query = (
+            sa.select(
+                raw_counts.c.bucket.label("bucket"),
+                sa.func.min(raw_counts.c.oldest).label("oldest"),
+                sa.func.sum(raw_counts.c.failed_count).label("failed_count"),
+                sa.func.sum(raw_counts.c.successful_count).label("successful_count"),
+            )
+            .select_from(raw_counts)
+            .group_by(raw_counts.c.bucket)
+            .order_by(sa.asc("oldest"))
+        )
+
+        result = await session.execute(query)
+
+    # Ensure that all buckets of time are present in the result even if no
+    # matching task runs occurred during the given time period.
+    buckets = [TaskRunCount(completed=0, failed=0) for _ in range(bucket_count)]
+
+    for row in result:
+        index = int(row.bucket)
+        buckets[index].completed = row.successful_count
+        buckets[index].failed = row.failed_count
+
+    return buckets


Test patch:
diff --git a/tests/server/api/ui/test_task_runs.py b/tests/server/api/ui/test_task_runs.py
new file mode 100644
--- /dev/null
+++ b/tests/server/api/ui/test_task_runs.py
@@ -0,0 +1,130 @@
+from typing import Tuple, cast
+
+import pendulum
+import pytest
+from httpx import AsyncClient
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from prefect.server import models
+from prefect.server.schemas import core, filters, states
+from prefect.server.utilities.schemas import DateTimeTZ
+from prefect.server.api.ui.task_runs import TaskRunCount
+
+
+class TestReadDashboardTaskRunCounts:
+    @pytest.fixture
+    def url(self) -> str:
+        return "/ui/task_runs/dashboard/counts"
+
+    @pytest.fixture
+    def time_window(self) -> Tuple[DateTimeTZ, DateTimeTZ]:
+        now = cast(DateTimeTZ, pendulum.datetime(2023, 6, 1, 18, tz="UTC"))
+        return (now, now.subtract(hours=8))
+
+    @pytest.fixture
+    def task_run_filter(self, time_window) -> filters.TaskRunFilter:
+        now, eight_hours_ago = time_window
+        return filters.TaskRunFilter(
+            start_time=filters.TaskRunFilterStartTime(
+                after_=eight_hours_ago, before_=now
+            )
+        )
+
+    @pytest.fixture
+    async def create_task_runs(
+        self,
+        session: AsyncSession,
+        flow_run,
+        time_window,
+    ):
+        now, eight_hours_ago = time_window
+        runs_to_create = 100
+        time_gap = (now - eight_hours_ago).as_timedelta() / runs_to_create
+
+        for i in range(runs_to_create):
+            if i % 2 == 0:
+                state_type = states.StateType.COMPLETED
+                state_name = "Completed"
+            else:
+                state_type = states.StateType.FAILED
+                state_name = "Failed"
+
+            await models.task_runs.create_task_run(
+                session=session,
+                task_run=core.TaskRun(
+                    flow_run_id=flow_run.id,
+                    task_key=f"task-{i}",
+                    dynamic_key=str(i),
+                    state_type=state_type,
+                    state_name=state_name,
+                    start_time=eight_hours_ago + (i * time_gap),
+                    end_time=eight_hours_ago + (i * time_gap),
+                ),
+            )
+
+        await session.commit()
+
+    async def test_requires_task_run_filter_start_time(
+        self,
+        url: str,
+        client: AsyncClient,
+    ):
+        task_run_filter = filters.TaskRunFilter()
+        response = await client.post(
+            url, json={"task_runs": task_run_filter.dict(json_compatible=True)}
+        )
+        assert response.status_code == 422
+        assert b"task_runs.start_time.after_ is required" in response.content
+
+    async def test_empty_state(
+        self,
+        url: str,
+        task_run_filter: filters.TaskRunFilter,
+        client: AsyncClient,
+    ):
+        response = await client.post(
+            url, json={"task_runs": task_run_filter.dict(json_compatible=True)}
+        )
+        assert response.status_code == 200
+
+        counts = [TaskRunCount(**count) for count in response.json()]
+        assert counts == [
+            TaskRunCount(completed=0, failed=0) for _ in range(len(counts))
+        ]
+
+    async def test_returns_completed_and_failed_counts(
+        self,
+        url: str,
+        task_run_filter: filters.TaskRunFilter,
+        client: AsyncClient,
+        create_task_runs,
+    ):
+        response = await client.post(
+            url, json={"task_runs": task_run_filter.dict(json_compatible=True)}
+        )
+        assert response.status_code == 200
+
+        counts = [TaskRunCount(**count) for count in response.json()]
+
+        assert counts == [
+            TaskRunCount(completed=3, failed=3),
+            TaskRunCount(completed=3, failed=2),
+            TaskRunCount(completed=2, failed=3),
+            TaskRunCount(completed=3, failed=2),
+            TaskRunCount(completed=2, failed=3),
+            TaskRunCount(completed=3, failed=2),
+            TaskRunCount(completed=2, failed=3),
+            TaskRunCount(completed=3, failed=2),
+            TaskRunCount(completed=2, failed=3),
+            TaskRunCount(completed=3, failed=2),
+            TaskRunCount(completed=2, failed=3),
+            TaskRunCount(completed=3, failed=2),
+            TaskRunCount(completed=2, failed=3),
+            TaskRunCount(completed=3, failed=2),
+            TaskRunCount(completed=2, failed=3),
+            TaskRunCount(completed=3, failed=2),
+            TaskRunCount(completed=2, failed=3),
+            TaskRunCount(completed=3, failed=2),
+            TaskRunCount(completed=2, failed=3),
+            TaskRunCount(completed=2, failed=2),
+        ]