Add work pool type to related resource (#16636)

PrefectHQ · Jan 8, 2025 · d321af1 · d321af1
1 parent e9890a9
commit d321af1
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 7 deletions.
diff --git a/src/prefect/server/models/events.py b/src/prefect/server/models/events.py
@@ -171,6 +171,7 @@ def _as_resource_data(
                 "name": work_pool.name,
                 "tags": [],
                 "role": "work-pool",
+                "type": work_pool.type,
             }
             if work_pool
             else {}
@@ -204,13 +205,16 @@ def _resource_data_as_related_resources(
         if kind in excluded_kinds or not data:
             continue
 
-        related.append(
-            {
-                "prefect.resource.id": f"prefect.{kind}.{data['id']}",
-                "prefect.resource.role": data["role"],
-                "prefect.resource.name": data["name"],
-            }
-        )
+        related_resource = {
+            "prefect.resource.id": f"prefect.{kind}.{data['id']}",
+            "prefect.resource.role": data["role"],
+            "prefect.resource.name": data["name"],
+        }
+
+        if kind == "work-pool":
+            related_resource["prefect.work-pool.type"] = data["type"]
+
+        related.append(related_resource)
 
     related += [
         {

diff --git a/tests/server/orchestration/test_flow_run_instrumentation_policies.py b/tests/server/orchestration/test_flow_run_instrumentation_policies.py
@@ -583,6 +583,7 @@ async def test_instrumenting_a_flow_run_from_a_work_queue(
                 "prefect.resource.id": f"prefect.work-pool.{work_pool.id}",
                 "prefect.resource.role": "work-pool",
                 "prefect.resource.name": work_pool.name,
+                "prefect.work-pool.type": work_pool.type,
             }
         )
         in event.related
@@ -880,6 +881,7 @@ async def test_caches_resource_data(
             "name": work_pool.name,
             "tags": [],
             "role": "work-pool",
+            "type": work_pool.type,
         },
         "task-run": {},
     }
@@ -960,6 +962,7 @@ async def test_caches_resource_data_for_subflow_runs(
             "name": work_pool.name,
             "tags": [],
             "role": "work-pool",
+            "type": work_pool.type,
         },
         "task-run": {
             "id": str(task_run.id),