Flow run state change events (#12825)

Author: chrisguidry

src/prefect/client/orchestration.py

@@ -3143,9 +3143,7 @@ async def read_automations(self) -> List[Automation]:
         response.raise_for_status()
         return pydantic.parse_obj_as(List[Automation], response.json())
 
-    async def find_automation(
-        self, id_or_name: str, exit_if_not_found: bool = True
-    ) -> Optional[Automation]:
+    async def find_automation(self, id_or_name: str) -> Optional[Automation]:
         try:
             id = UUID(id_or_name)
         except ValueError:

src/prefect/server/api/deployments.py

@@ -115,7 +115,7 @@ async def create_deployment(
         elif deployment.work_queue_name:
             # If just a queue name was provided, ensure that the queue exists and
             # get its ID.
-            work_queue = await models.work_queues._ensure_work_queue_exists(
+            work_queue = await models.work_queues.ensure_work_queue_exists(
                 session=session, name=deployment.work_queue_name
             )
             deployment_dict["work_queue_id"] = work_queue.id

src/prefect/server/models/deployments.py

@@ -218,8 +218,8 @@ async def update_deployment(
     elif deployment.work_queue_name:
         # If just a queue name was provided, ensure the queue exists and
         # get its ID.
-        work_queue = await models.work_queues._ensure_work_queue_exists(
-            session=session, name=update_data["work_queue_name"], db=db
+        work_queue = await models.work_queues.ensure_work_queue_exists(
+            session=session, name=update_data["work_queue_name"]
         )
         update_data["work_queue_id"] = work_queue.id
 

src/prefect/server/models/events.py

@@ -0,0 +1,287 @@
+from datetime import datetime, timedelta
+from typing import Any, Dict, List, MutableMapping, Optional, Set, Union
+from uuid import UUID
+
+from cachetools import TTLCache
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from prefect.server import models, schemas
+from prefect.server.database.orm_models import (
+    ORMDeployment,
+    ORMFlow,
+    ORMFlowRun,
+    ORMFlowRunState,
+    ORMTaskRun,
+    ORMTaskRunState,
+    ORMWorkPool,
+    ORMWorkQueue,
+)
+from prefect.server.events.schemas.events import Event
+from prefect.server.models import deployments
+from prefect.settings import PREFECT_API_EVENTS_RELATED_RESOURCE_CACHE_TTL
+from prefect.utilities.text import truncated_to
+
+ResourceData = Dict[str, Dict[str, Any]]
+RelatedResourceList = List[Dict[str, str]]
+
+
+# Some users use state messages to convey error messages and large results; let's
+# truncate them so they don't blow out the size of a message
+TRUNCATE_STATE_MESSAGES_AT = 100_000
+
+
+_flow_run_resource_data_cache: MutableMapping[UUID, ResourceData] = TTLCache(
+    maxsize=1000,
+    ttl=PREFECT_API_EVENTS_RELATED_RESOURCE_CACHE_TTL.value().total_seconds(),
+)
+
+
+async def flow_run_state_change_event(
+    session: AsyncSession,
+    occurred: datetime,
+    flow_run: ORMFlowRun,
+    initial_state_id: Optional[UUID],
+    initial_state: Optional[schemas.states.State],
+    validated_state_id: Optional[UUID],
+    validated_state: schemas.states.State,
+) -> Event:
+    return Event(
+        occurred=occurred,
+        event=f"prefect.flow-run.{validated_state.name}",
+        resource={
+            "prefect.resource.id": f"prefect.flow-run.{flow_run.id}",
+            "prefect.resource.name": flow_run.name,
+            "prefect.state-message": truncated_to(
+                TRUNCATE_STATE_MESSAGES_AT, validated_state.message
+            ),
+            "prefect.state-name": validated_state.name or "",
+            "prefect.state-timestamp": (
+                validated_state.timestamp.isoformat()
+                if validated_state.timestamp
+                else None
+            ),
+            "prefect.state-type": validated_state.type.value,
+        },
+        related=await _flow_run_related_resources_from_orm(
+            session=session, flow_run=flow_run
+        ),
+        payload={
+            "intended": {
+                "from": _state_type(initial_state),
+                "to": _state_type(validated_state),
+            },
+            "initial_state": state_payload(initial_state),
+            "validated_state": state_payload(validated_state),
+        },
+        # Here we use the state's ID as the ID of the event as well, in order to
+        # establish the ordering of the state-change events
+        id=validated_state_id,
+        follows=initial_state_id if _timing_is_tight(occurred, initial_state) else None,
+    )
+
+
+async def _flow_run_related_resources_from_orm(
+    session: AsyncSession, flow_run: ORMFlowRun
+) -> RelatedResourceList:
+    resource_data = _flow_run_resource_data_cache.get(flow_run.id)
+    if not resource_data:
+        flow = await models.flows.read_flow(session=session, flow_id=flow_run.flow_id)
+        deployment: Optional[ORMDeployment] = None
+        if flow_run.deployment_id:
+            deployment = await deployments.read_deployment(
+                session, deployment_id=flow_run.deployment_id
+            )
+
+        work_queue = None
+        if flow_run.work_queue_id:
+            work_queue = await models.work_queues.read_work_queue(
+                session, work_queue_id=flow_run.work_queue_id
+            )
+
+        work_pool = work_queue.work_pool if work_queue is not None else None
+
+        task_run: Optional[ORMTaskRun] = None
+        if flow_run.parent_task_run_id:
+            task_run = await models.task_runs.read_task_run(
+                session,
+                task_run_id=flow_run.parent_task_run_id,
+            )
+
+        print(task_run)
+        if task_run:
+            print("events", task_run.id, task_run.tags)
+
+        resource_data = _as_resource_data(
+            flow_run, flow, deployment, work_queue, work_pool, task_run
+        )
+        _flow_run_resource_data_cache[flow_run.id] = resource_data
+
+    return _resource_data_as_related_resources(
+        resource_data,
+        excluded_kinds=["flow-run"],
+    ) + _provenance_as_related_resources(flow_run.created_by)
+
+
+def _as_resource_data(
+    flow_run: ORMFlowRun,
+    flow: Union[ORMFlow, schemas.core.Flow, None],
+    deployment: Union[ORMDeployment, schemas.responses.DeploymentResponse, None],
+    work_queue: Union[ORMWorkQueue, schemas.responses.WorkQueueResponse, None],
+    work_pool: Union[ORMWorkPool, schemas.core.WorkPool, None],
+    task_run: Union[ORMTaskRun, schemas.core.TaskRun, None] = None,
+) -> ResourceData:
+    return {
+        "flow-run": {
+            "id": str(flow_run.id),
+            "name": flow_run.name,
+            "tags": flow_run.tags if flow_run.tags else [],
+            "role": "flow-run",
+        },
+        "flow": (
+            {
+                "id": str(flow.id),
+                "name": flow.name,
+                "tags": flow.tags if flow.tags else [],
+                "role": "flow",
+            }
+            if flow
+            else {}
+        ),
+        "deployment": (
+            {
+                "id": str(deployment.id),
+                "name": deployment.name,
+                "tags": deployment.tags if deployment.tags else [],
+                "role": "deployment",
+            }
+            if deployment
+            else {}
+        ),
+        "work-queue": (
+            {
+                "id": str(work_queue.id),
+                "name": work_queue.name,
+                "tags": [],
+                "role": "work-queue",
+            }
+            if work_queue
+            else {}
+        ),
+        "work-pool": (
+            {
+                "id": str(work_pool.id),
+                "name": work_pool.name,
+                "tags": [],
+                "role": "work-pool",
+            }
+            if work_pool
+            else {}
+        ),
+        "task-run": (
+            {
+                "id": str(task_run.id),
+                "name": task_run.name,
+                "tags": task_run.tags if task_run.tags else [],
+                "role": "task-run",
+            }
+            if task_run
+            else {}
+        ),
+    }
+
+
+def _resource_data_as_related_resources(
+    resource_data: ResourceData,
+    excluded_kinds: Optional[List[str]] = None,
+) -> RelatedResourceList:
+    related = []
+    tags: Set[str] = set()
+
+    if excluded_kinds is None:
+        excluded_kinds = []
+
+    for kind, data in resource_data.items():
+        tags |= set(data.get("tags", []))
+
+        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 += [
+        {
+            "prefect.resource.id": f"prefect.tag.{tag}",
+            "prefect.resource.role": "tag",
+        }
+        for tag in sorted(tags)
+    ]
+
+    return related
+
+
+def _provenance_as_related_resources(
+    created_by: Optional[schemas.core.CreatedBy],
+) -> RelatedResourceList:
+    if not created_by:
+        return []
+
+    resource_id: str
+
+    if created_by.type == "DEPLOYMENT":
+        resource_id = f"prefect.deployment.{created_by.id}"
+    elif created_by.type == "AUTOMATION":
+        resource_id = f"prefect.automation.{created_by.id}"
+    else:
+        return []
+
+    related = {
+        "prefect.resource.id": resource_id,
+        "prefect.resource.role": "creator",
+    }
+    if created_by.display_value:
+        related["prefect.resource.name"] = created_by.display_value
+    return [related]
+
+
+def _state_type(
+    state: Union[ORMFlowRunState, ORMTaskRunState, Optional[schemas.states.State]],
+) -> Optional[str]:
+    return str(state.type.value) if state else None
+
+
+def state_payload(state: Optional[schemas.states.State]) -> Optional[Dict[str, str]]:
+    """Given a State, return the essential string parts of it for use in an
+    event payload"""
+    if not state:
+        return None
+    payload: Dict[str, str] = {"type": state.type.value}
+    if state.name:
+        payload["name"] = state.name
+    if state.message:
+        payload["message"] = truncated_to(TRUNCATE_STATE_MESSAGES_AT, state.message)
+    if state.is_paused():
+        payload["pause_reschedule"] = str(state.state_details.pause_reschedule).lower()
+    return payload
+
+
+def _timing_is_tight(
+    occurred: datetime,
+    initial_state: Union[
+        ORMFlowRunState, ORMTaskRunState, Optional[schemas.states.State]
+    ],
+) -> bool:
+    # Only connect events with event.follows if the timing here is tight, which will
+    # help us resolve the order of these events if they happen to be delivered out of
+    # order.  If the preceding state change happened a while back, don't worry about
+    # it because the order is very likely to be unambiguous.
+    TIGHT_TIMING = timedelta(minutes=5)
+    if initial_state and initial_state.timestamp:
+        return bool(-TIGHT_TIMING < (occurred - initial_state.timestamp) < TIGHT_TIMING)
+
+    return False