Index task runs to improve GET performance

Author: daniel-slaugh

domains/etl/migrations/0006_taskrun_latest_lookup_index.py

@@ -0,0 +1,20 @@
+# Generated by Django 5.2.2 on 2026-03-13 00:00
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("etl", "0005_fix_celery_task_names"),
+    ]
+
+    operations = [
+        migrations.AddIndex(
+            model_name="taskrun",
+            index=models.Index(
+                fields=["task", "-started_at"],
+                name="etl_taskrun_task_started_idx",
+            ),
+        ),
+    ]

domains/etl/models/run.py

@@ -59,3 +59,11 @@ def message(self) -> str | None:
     @property
     def failure_count(self) -> int | None:
         return self.extract_failure_count(self.result)
+
+    class Meta:
+        indexes = [
+            models.Index(
+                fields=["task", "-started_at"],
+                name="etl_taskrun_task_started_idx",
+            ),
+        ]

domains/etl/services/task.py

@@ -139,6 +139,22 @@ def build_task_response(
         include_latest_run_result: bool = True,
         include_data_connection_settings: bool = True,
     ) -> dict:
+        latest_task_run = getattr(task, "latest_task_run", None)
+        latest_run_message = None
+        latest_run_failure_count = None
+
+        if latest_task_run:
+            latest_run_message = getattr(latest_task_run, "message_text", None)
+            latest_run_failure_count = getattr(
+                latest_task_run, "failure_count_value", None
+            )
+
+            if include_latest_run_result:
+                if latest_run_message is None:
+                    latest_run_message = latest_task_run.message
+                if latest_run_failure_count is None:
+                    latest_run_failure_count = latest_task_run.failure_count
+
         response = {
             "id": task.id,
             "name": task.name,
@@ -155,18 +171,33 @@ def build_task_response(
                 "intervalPeriod": task.periodic_task.interval.period if task.periodic_task.interval else None,
             } if task.periodic_task else None,
             "latest_run": {
-                "id": getattr(task, "latest_run_id", None),
-                "status": getattr(task, "latest_run_status", None),
-                "message": getattr(task, "latest_run_message", None),
-                "failure_count": getattr(task, "latest_run_failure_count", None),
+                "id": latest_task_run.id if latest_task_run else getattr(task, "latest_run_id", None),
+                "status": latest_task_run.status if latest_task_run else getattr(task, "latest_run_status", None),
+                "message": latest_run_message if latest_task_run else getattr(task, "latest_run_message", None),
+                "failure_count": (
+                    latest_run_failure_count
+                    if latest_task_run
+                    else getattr(task, "latest_run_failure_count", None)
+                ),
                 "result": (
-                    getattr(task, "latest_run_result", None)
-                    if include_latest_run_result
-                    else None
+                    (
+                        latest_task_run.result
+                        if latest_task_run
+                        else getattr(task, "latest_run_result", None)
+                    )
+                    if include_latest_run_result else None
+                ),
+                "started_at": (
+                    latest_task_run.started_at
+                    if latest_task_run
+                    else getattr(task, "latest_run_started_at", None)
                 ),
-                "started_at": getattr(task, "latest_run_started_at", None),
-                "finished_at": getattr(task, "latest_run_finished_at", None),
-            } if getattr(task, "latest_run_id", None) else None,
+                "finished_at": (
+                    latest_task_run.finished_at
+                    if latest_task_run
+                    else getattr(task, "latest_run_finished_at", None)
+                ),
+            } if latest_task_run or getattr(task, "latest_run_id", None) else None,
             "extractor_variables": task.extractor_variables,
             "transformer_variables": task.transformer_variables,
             "loader_variables": task.loader_variables,
@@ -265,6 +296,72 @@ def annotate_target_identifiers(queryset: QuerySet) -> QuerySet:
             )
         )
 
+    @staticmethod
+    def annotate_latest_task_run_fields(queryset: QuerySet) -> QuerySet:
+        task_result_queryset = (
+            TaskRun.objects
+            .filter(task_id=OuterRef("pk"))
+            .order_by("-started_at", "-id")
+        )
+        return queryset.annotate(
+            latest_run_id=Subquery(
+                task_result_queryset.values("id")[:1]
+            ),
+            latest_run_status=Subquery(
+                task_result_queryset.values("status")[:1]
+            ),
+            latest_run_started_at=Subquery(
+                task_result_queryset.values("started_at")[:1]
+            ),
+            latest_run_finished_at=Subquery(
+                task_result_queryset.values("finished_at")[:1]
+            ),
+        )
+
+    @staticmethod
+    def get_latest_runs_for_tasks(
+        task_ids: list[uuid.UUID],
+        include_result: bool = True,
+    ) -> dict[uuid.UUID, TaskRun]:
+        if not task_ids:
+            return {}
+
+        latest_runs = (
+            TaskRun.objects
+            .filter(task_id__in=task_ids)
+            .annotate(
+                message_text=Coalesce(
+                    KeyTextTransform("message", "result"),
+                    KeyTextTransform("summary", "result"),
+                    KeyTextTransform("statusMessage", "result"),
+                    KeyTextTransform("status_message", "result"),
+                    KeyTextTransform("failureReason", "result"),
+                    KeyTextTransform("failure_reason", "result"),
+                    KeyTextTransform("error", "result"),
+                ),
+                failure_count_value=Coalesce(
+                    Cast(
+                        KeyTextTransform("failure_count", "result"),
+                        IntegerField(),
+                    ),
+                    Cast(
+                        KeyTextTransform("failureCount", "result"),
+                        IntegerField(),
+                    ),
+                    output_field=IntegerField(),
+                ),
+            )
+            .order_by("task_id", "-started_at", "-id")
+        )
+
+        if not include_result:
+            latest_runs = latest_runs.defer("result")
+
+        return {
+            task_run.task_id: task_run
+            for task_run in latest_runs.distinct("task_id")
+        }
+
     @staticmethod
     def annotate_latest_task_result(
         queryset: QuerySet,
@@ -273,7 +370,7 @@ def annotate_latest_task_result(
         task_result_queryset = (
             TaskRun.objects
             .filter(task_id=OuterRef("pk"))
-            .order_by("-started_at")
+            .order_by("-started_at", "-id")
         )
         annotations = {
             "latest_run_id": Subquery(
@@ -349,12 +446,34 @@ def list(
             if include_data_connection_settings is None
             else include_data_connection_settings
         )
+        filtering = filtering or {}
+        order_by = order_by or []
         queryset = Task.objects
 
-        queryset = self.annotate_latest_task_result(
-            queryset,
-            include_result=include_latest_run_result,
-        )
+        if (
+            any(
+                field in filtering
+                for field in [
+                    "latest_run_status",
+                    "latest_run_started_at__lte",
+                    "latest_run_started_at__gte",
+                    "latest_run_finished_at__lte",
+                    "latest_run_finished_at__gte",
+                ]
+            )
+            or any(
+                field.lstrip("-")
+                in {
+                    "latestRunStatus",
+                    "latestRunStartedAt",
+                    "latestRunFinishedAt",
+                }
+                for field in order_by
+            )
+        ):
+            queryset = self.annotate_latest_task_run_fields(
+                queryset
+            )
 
         for field in [
             "workspace_id",
@@ -425,14 +544,22 @@ def list(
         queryset = queryset.visible(principal=principal).distinct()  # noqa
         queryset, count = self.apply_pagination(queryset, response, page, page_size)
 
+        queryset = list(queryset.all())
+        latest_runs_by_task_id = self.get_latest_runs_for_tasks(
+            [task.id for task in queryset],
+            include_result=include_latest_run_result,
+        )
+        for task in queryset:
+            task.latest_task_run = latest_runs_by_task_id.get(task.id)
+
         return [
             self.build_task_response(
                 task,
                 expand=expand_related,
                 include_mappings=include_mappings,
                 include_latest_run_result=include_latest_run_result,
                 include_data_connection_settings=include_data_connection_settings,
-            ) for task in queryset.all()
+            ) for task in queryset
         ]
 
     def get(

tests/etl/services/test_task.py

@@ -1,5 +1,6 @@
 import pytest
 import uuid
+from datetime import datetime
 from collections import Counter
 from ninja.errors import HttpError
 from django.http import HttpResponse
@@ -233,6 +234,57 @@ def test_list_task_can_expose_failure_count_without_latest_run_result(get_princi
     assert lean_task["latest_run"]["result"] is None
 
 
+def test_list_task_can_filter_and_order_by_latest_run_fields(get_principal):
+    task = Task.objects.create(
+        name="Later Task",
+        workspace_id=uuid.UUID("b27c51a0-7374-462d-8a53-d97d47176c10"),
+        data_connection_id=uuid.UUID("019adb5c-da8b-7970-877d-c3b4ca37cc60"),
+        orchestration_system_id=uuid.UUID("7cb900d2-eb11-4a59-a05b-dd02d95af312"),
+        extractor_variables={},
+        transformer_variables={},
+        loader_variables={},
+    )
+    latest_run = TaskRun.objects.create(
+        task=task,
+        status="FAILURE",
+        result={
+            "summary": "Later failed",
+            "failure_count": 3,
+        },
+    )
+    TaskRun.objects.filter(pk=latest_run.pk).update(
+        started_at=timezone.make_aware(datetime(2025, 1, 3, 1, 0, 0)),
+        finished_at=timezone.make_aware(datetime(2025, 1, 3, 2, 0, 0)),
+    )
+
+    filtered = task_service.list(
+        principal=get_principal("owner"),
+        response=HttpResponse(),
+        page=1,
+        page_size=100,
+        order_by=["-latestRunStartedAt"],
+        filtering={"latest_run_status": ["FAILURE"]},
+        include_latest_run_result=False,
+    )
+
+    assert [task["name"] for task in filtered] == ["Later Task"]
+    assert filtered[0]["latest_run"]["status"] == "FAILURE"
+    assert filtered[0]["latest_run"]["message"] == "Later failed"
+    assert filtered[0]["latest_run"]["failure_count"] == 3
+    assert filtered[0]["latest_run"]["result"] is None
+
+    ordered = task_service.list(
+        principal=get_principal("owner"),
+        response=HttpResponse(),
+        page=1,
+        page_size=100,
+        order_by=["-latestRunStartedAt"],
+        filtering={},
+    )
+
+    assert [task["name"] for task in ordered[:2]] == ["Later Task", "Test ETL Task"]
+
+
 def test_run_task_returns_a_new_running_run(get_principal, monkeypatch, settings):
     settings.CELERY_ENABLED = True
     principal = get_principal("owner")