Repository: PrefectHQ/prefect

Issue:
Include deployment, work-queue and work-pool as related resources when emitting events.
### First check

- [X] I added a descriptive title to this issue.
- [X] I used the GitHub search to find a similar request and didn't find it.
- [X] I searched the Prefect documentation for this feature.

### Prefect Version

2.x

### Describe the current behavior

Currently when an event is emitted prefect will only pull in the related flow and flow run.

### Describe the proposed behavior

When emitting events the event should include flow, flow run, deployment, work queue, and work pool as related resources.

### Example Use

_No response_

### Additional context

_No response_


Discussion:

Resolving patch:
diff --git a/src/prefect/events/related.py b/src/prefect/events/related.py
--- a/src/prefect/events/related.py
+++ b/src/prefect/events/related.py
@@ -1,15 +1,31 @@
-from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Union
+import asyncio
+import pendulum
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    Union,
+)
 from uuid import UUID
+from pendulum.datetime import DateTime
 
 from .schemas import RelatedResource
 
 if TYPE_CHECKING:
-    from prefect.client.schemas import FlowRun
-    from prefect.server.schemas.core import Flow
+    from prefect.server.utilities.schemas import ORMBaseModel
 
+ResourceCacheEntry = Dict[str, Union[str, "ORMBaseModel", None]]
+RelatedResourceCache = Dict[str, Tuple[ResourceCacheEntry, DateTime]]
 
-ObjectDict = Dict[str, Union["Flow", "FlowRun"]]
-related_resource_cache: Dict[UUID, ObjectDict] = {}
+MAX_CACHE_SIZE = 100
+RESOURCE_CACHE: RelatedResourceCache = {}
 
 
 def tags_as_related_resources(tags: Iterable[str]) -> List[RelatedResource]:
@@ -57,32 +73,91 @@ async def related_resources_from_run_context(
         else task_run_context.task_run.flow_run_id
     )
 
-    objects = related_resource_cache.get(flow_run_id, {})
+    related_objects: list[ResourceCacheEntry] = []
 
     async with get_client() as client:
-        if "flow-run" not in objects:
-            flow_run = await client.read_flow_run(flow_run_id)
-            if flow_run:
-                objects["flow-run"] = flow_run
 
-        if "flow" not in objects and "flow-run" in objects:
-            flow = await client.read_flow(objects["flow-run"].flow_id)
-            if flow:
-                objects["flow"] = flow
-
-        related_resource_cache[flow_run_id] = objects
+        async def dummy_read():
+            return {}
+
+        related_objects = [
+            await _get_and_cache_related_object(
+                kind="flow-run",
+                role="flow-run",
+                client_method=client.read_flow_run,
+                obj_id=flow_run_id,
+                cache=RESOURCE_CACHE,
+            )
+        ]
+
+        flow_run = related_objects[0]["object"]
+
+        if flow_run:
+            related_objects += list(
+                await asyncio.gather(
+                    _get_and_cache_related_object(
+                        kind="flow",
+                        role="flow",
+                        client_method=client.read_flow,
+                        obj_id=flow_run.flow_id,
+                        cache=RESOURCE_CACHE,
+                    ),
+                    (
+                        _get_and_cache_related_object(
+                            kind="deployment",
+                            role="deployment",
+                            client_method=client.read_deployment,
+                            obj_id=flow_run.deployment_id,
+                            cache=RESOURCE_CACHE,
+                        )
+                        if flow_run.deployment_id
+                        else dummy_read()
+                    ),
+                    (
+                        _get_and_cache_related_object(
+                            kind="work-queue",
+                            role="work-queue",
+                            client_method=client.read_work_queue,
+                            obj_id=flow_run.work_queue_id,
+                            cache=RESOURCE_CACHE,
+                        )
+                        if flow_run.work_queue_id
+                        else dummy_read()
+                    ),
+                    (
+                        _get_and_cache_related_object(
+                            kind="work-pool",
+                            role="work-pool",
+                            client_method=client.read_work_pool,
+                            obj_id=flow_run.work_pool_name,
+                            cache=RESOURCE_CACHE,
+                        )
+                        if flow_run.work_pool_name
+                        else dummy_read()
+                    ),
+                )
+            )
 
     related = []
     tags = set()
 
-    for kind, obj in objects.items():
-        resource = object_as_related_resource(kind=kind, role=kind, object=obj)
+    for entry in related_objects:
+        obj_ = entry.get("object")
+        if obj_ is None:
+            continue
+
+        assert isinstance(entry["kind"], str) and isinstance(entry["role"], str)
+
+        resource = object_as_related_resource(
+            kind=entry["kind"], role=entry["kind"], object=obj_
+        )
 
         if resource.id in exclude:
             continue
 
         related.append(resource)
-        tags |= set(obj.tags)
+        if hasattr(obj_, "tags"):
+            tags |= set(obj_.tags)
 
     related += [
         resource
@@ -91,3 +166,43 @@ async def related_resources_from_run_context(
     ]
 
     return related
+
+
+async def _get_and_cache_related_object(
+    kind: str,
+    role: str,
+    client_method: Callable[[Union[UUID, str]], Awaitable[Optional["ORMBaseModel"]]],
+    obj_id: Union[UUID, str],
+    cache: RelatedResourceCache,
+) -> ResourceCacheEntry:
+    cache_key = f"{kind}.{obj_id}"
+    entry = None
+
+    if cache_key in cache:
+        entry, _ = cache[cache_key]
+    else:
+        obj_ = await client_method(obj_id)
+        entry = {
+            "kind": kind,
+            "object": obj_,
+        }
+
+    cache[cache_key] = (entry, pendulum.now("UTC"))
+
+    # In the case of a worker or agent this cache could be long-lived. To keep
+    # from running out of memory only keep `MAX_CACHE_SIZE` entries in the
+    # cache.
+    if len(cache) > MAX_CACHE_SIZE:
+        oldest_key = sorted(
+            [(key, timestamp) for key, (_, timestamp) in cache.items()],
+            key=lambda k: k[1],
+        )[0][0]
+
+        if oldest_key:
+            del cache[oldest_key]
+
+    # Because the role is event specific and can change depending on the
+    # type of event being emitted, this adds the role from the args to the
+    # entry before returning it rather than storing it in the cache.
+    entry["role"] = role
+    return entry
diff --git a/src/prefect/server/schemas/responses.py b/src/prefect/server/schemas/responses.py
--- a/src/prefect/server/schemas/responses.py
+++ b/src/prefect/server/schemas/responses.py
@@ -165,6 +165,7 @@ class FlowRunResponse(ORMBaseModel):
     flow_id: UUID = FieldFrom(schemas.core.FlowRun)
     state_id: Optional[UUID] = FieldFrom(schemas.core.FlowRun)
     deployment_id: Optional[UUID] = FieldFrom(schemas.core.FlowRun)
+    work_queue_id: Optional[UUID] = FieldFrom(schemas.core.FlowRun)
     work_queue_name: Optional[str] = FieldFrom(schemas.core.FlowRun)
     flow_version: Optional[str] = FieldFrom(schemas.core.FlowRun)
     parameters: dict = FieldFrom(schemas.core.FlowRun)
@@ -187,6 +188,10 @@ class FlowRunResponse(ORMBaseModel):
     infrastructure_document_id: Optional[UUID] = FieldFrom(schemas.core.FlowRun)
     infrastructure_pid: Optional[str] = FieldFrom(schemas.core.FlowRun)
     created_by: Optional[CreatedBy] = FieldFrom(schemas.core.FlowRun)
+    work_pool_id: Optional[UUID] = Field(
+        default=None,
+        description="The id of the flow run's work pool.",
+    )
     work_pool_name: Optional[str] = Field(
         default=None,
         description="The name of the flow run's work pool.",
@@ -198,8 +203,10 @@ class FlowRunResponse(ORMBaseModel):
     def from_orm(cls, orm_flow_run: "prefect.server.database.orm_models.ORMFlowRun"):
         response = super().from_orm(orm_flow_run)
         if orm_flow_run.work_queue:
+            response.work_queue_id = orm_flow_run.work_queue.id
             response.work_queue_name = orm_flow_run.work_queue.name
             if orm_flow_run.work_queue.work_pool:
+                response.work_pool_id = orm_flow_run.work_queue.work_pool.id
                 response.work_pool_name = orm_flow_run.work_queue.work_pool.name
 
         return response


Test patch:
diff --git a/tests/events/test_events_related_from_context.py b/tests/events/test_events_related_from_context.py
--- a/tests/events/test_events_related_from_context.py
+++ b/tests/events/test_events_related_from_context.py
@@ -3,23 +3,33 @@
 
 import pytest
 
-from prefect import flow, tags, task
+from prefect import flow, task
 from prefect.client.orchestration import get_client
 from prefect.context import FlowRunContext, TaskRunContext
-from prefect.events.related import related_resources_from_run_context
+from prefect.events.related import (
+    MAX_CACHE_SIZE,
+    related_resources_from_run_context,
+    _get_and_cache_related_object,
+)
 from prefect.events.schemas import RelatedResource
+from prefect.states import Running
 
 
 @pytest.fixture
 async def spy_client(test_database_connection_url):
     async with get_client() as client:
         exit_stack = ExitStack()
-        exit_stack.enter_context(
-            mock.patch.object(client, "read_flow", wraps=client.read_flow),
-        )
-        exit_stack.enter_context(
-            mock.patch.object(client, "read_flow_run", wraps=client.read_flow_run),
-        )
+
+        for method in [
+            "read_flow",
+            "read_flow_run",
+            "read_deployment",
+            "read_work_queue",
+            "read_work_pool",
+        ]:
+            exit_stack.enter_context(
+                mock.patch.object(client, method, wraps=getattr(client, method)),
+            )
 
         class NoOpClientWrapper:
             def __init__(self, client):
@@ -41,19 +51,21 @@ async def test_gracefully_handles_missing_context():
     assert related == []
 
 
-async def test_gets_related_from_run_context(orion_client):
-    @flow
-    async def test_flow():
-        return await related_resources_from_run_context()
+async def test_gets_related_from_run_context(
+    orion_client, work_queue_1, worker_deployment_wq1
+):
+    flow_run = await orion_client.create_flow_run_from_deployment(
+        worker_deployment_wq1.id,
+        state=Running(),
+        tags=["flow-run-one"],
+    )
 
-    with tags("testing"):
-        state = await test_flow._run()
+    with FlowRunContext.construct(flow_run=flow_run):
+        related = await related_resources_from_run_context()
 
-    flow_run = await orion_client.read_flow_run(state.state_details.flow_run_id)
+    work_pool = work_queue_1.work_pool
     db_flow = await orion_client.read_flow(flow_run.flow_id)
 
-    related = await state.result()
-
     assert related == [
         RelatedResource(
             __root__={
@@ -71,7 +83,34 @@ async def test_flow():
         ),
         RelatedResource(
             __root__={
-                "prefect.resource.id": "prefect.tag.testing",
+                "prefect.resource.id": f"prefect.deployment.{worker_deployment_wq1.id}",
+                "prefect.resource.role": "deployment",
+                "prefect.resource.name": worker_deployment_wq1.name,
+            }
+        ),
+        RelatedResource(
+            __root__={
+                "prefect.resource.id": f"prefect.work-queue.{work_queue_1.id}",
+                "prefect.resource.role": "work-queue",
+                "prefect.resource.name": work_queue_1.name,
+            }
+        ),
+        RelatedResource(
+            __root__={
+                "prefect.resource.id": f"prefect.work-pool.{work_pool.id}",
+                "prefect.resource.role": "work-pool",
+                "prefect.resource.name": work_pool.name,
+            }
+        ),
+        RelatedResource(
+            __root__={
+                "prefect.resource.id": "prefect.tag.flow-run-one",
+                "prefect.resource.role": "tag",
+            }
+        ),
+        RelatedResource(
+            __root__={
+                "prefect.resource.id": "prefect.tag.test",
                 "prefect.resource.role": "tag",
             }
         ),
@@ -165,3 +204,38 @@ async def test_flow():
 
     spy_client.client.read_flow_run.assert_called_once()
     spy_client.client.read_flow.assert_called_once()
+
+
+async def test_lru_cache_evicts_oldest():
+    cache = {}
+
+    async def fetch(obj_id):
+        return obj_id
+
+    await _get_and_cache_related_object("flow-run", "flow-run", fetch, "👴", cache)
+    assert "flow-run.👴" in cache
+
+    await _get_and_cache_related_object("flow-run", "flow-run", fetch, "👩", cache)
+    assert "flow-run.👴" in cache
+
+    for i in range(MAX_CACHE_SIZE):
+        await _get_and_cache_related_object(
+            "flow-run", "flow-run", fetch, f"👶 {i}", cache
+        )
+
+    assert "flow-run.👴" not in cache
+
+
+async def test_lru_cache_timestamp_updated():
+    cache = {}
+
+    async def fetch(obj_id):
+        return obj_id
+
+    await _get_and_cache_related_object("flow-run", "flow-run", fetch, "👴", cache)
+    _, timestamp = cache["flow-run.👴"]
+
+    await _get_and_cache_related_object("flow-run", "flow-run", fetch, "👴", cache)
+    _, next_timestamp = cache["flow-run.👴"]
+
+    assert next_timestamp > timestamp