async def state_event(flow_run_id): flow_run = await models.FlowRun.where(id=flow_run_id).first({ "id": True, "name": True, "version": True, "flow": {"id", "name", "version_group_id"}, "tenant": {"id", "slug"}, "current_state": {"state", "serialized_state", "version"}, }) return events.FlowRunStateChange( flow_run=flow_run, tenant=flow_run.tenant, state=flow_run.current_state, flow=flow_run.flow, )
async def set_flow_run_state(flow_run_id: str, state: State, version: int = None, agent_id: str = None) -> models.FlowRunState: """ Updates a flow run state. Args: - flow_run_id (str): the flow run id to update - state (State): the new state - version (int): a version to enforce version-locking - agent_id (str): the ID of an agent instance setting the state Returns: - models.FlowRunState """ if flow_run_id is None: raise ValueError(f"Invalid flow run ID.") where = { "id": { "_eq": flow_run_id }, "_or": [ # EITHER version locking is enabled and versions match { "version": { "_eq": version }, "flow": { "flow_group": { "settings": { "_contains": { "version_locking_enabled": True } } } }, }, # OR version locking is not enabled { "flow": { "flow_group": { "_not": { "settings": { "_contains": { "version_locking_enabled": True } } } } } }, ], } flow_run = await models.FlowRun.where(where).first({ "id": True, "state": True, "name": True, "version": True, "flow": {"id", "name", "flow_group_id", "version_group_id"}, "tenant": {"id", "slug"}, }) if not flow_run: raise ValueError(f"State update failed for flow run ID {flow_run_id}") # -------------------------------------------------------- # apply downstream updates # -------------------------------------------------------- # FOR CANCELLED STATES: # - set all non-finished task run states to Cancelled if isinstance(state, Cancelled): task_runs = await models.TaskRun.where({ "flow_run_id": { "_eq": flow_run_id } }).get({"id", "serialized_state"}) to_cancel = [ t for t in task_runs if not state_schema.load(t.serialized_state).is_finished() ] # For a run with many tasks this may be a lot of tasks - at some point # we might want to batch this rather than kicking off lots of asyncio # tasks at once. await asyncio.gather( *(api.states.set_task_run_state(t.id, state) for t in to_cancel), return_exceptions=True, ) # -------------------------------------------------------- # insert the new state in the database # -------------------------------------------------------- flow_run_state = models.FlowRunState( id=str(uuid.uuid4()), tenant_id=flow_run.tenant_id, flow_run_id=flow_run_id, version=(flow_run.version or 0) + 1, state=type(state).__name__, timestamp=pendulum.now("UTC"), message=state.message, result=state.result, start_time=getattr(state, "start_time", None), serialized_state=state.serialize(), ) await flow_run_state.insert() # -------------------------------------------------------- # apply downstream updates # -------------------------------------------------------- # FOR RUNNING STATES: # - update the flow run heartbeat if state.is_running() or state.is_submitted(): await api.runs.update_flow_run_heartbeat(flow_run_id=flow_run_id) # Set agent ID on flow run when submitted by agent if state.is_submitted() and agent_id: await api.runs.update_flow_run_agent(flow_run_id=flow_run_id, agent_id=agent_id) # -------------------------------------------------------- # call cloud hooks # -------------------------------------------------------- event = events.FlowRunStateChange( flow_run=flow_run, state=flow_run_state, flow=flow_run.flow, tenant=flow_run.tenant, ) asyncio.create_task(api.cloud_hooks.call_hooks(event)) return flow_run_state
async def test_cloud_hook( cloud_hook_id: str, flow_run_id: str = None, state: prefect.engine.state.State = None, ): """ Sends a test payload to a hook Args: - cloud_hook_id (str): the hook id - flow_run_id (str, optional): a flow run ID to test with. If provided, the event will pull event details from the flow run - state (State, optional): a sample state """ if state is None: state = prefect.engine.state.Success(message="Test success state") serialized_state = state.serialize() if cloud_hook_id is None: raise ValueError("Invalid ID") # verify that the hook actually exists hook = await models.CloudHook.where(id=cloud_hook_id).first( {"type", "config", "tenant_id"}) if not hook: raise ValueError("Invalid ID") # if a flow run has been provided, construct an event mirroring the actual event # we'd expect to see if flow_run_id: flow_run = await models.FlowRun.where(id=flow_run_id).first( {"id", "tenant_id", "flow_id", "name", "version"}) new_flow_run = Box( id=flow_run.id, is_test_event=True, tenant_id=flow_run.tenant_id, flow_id=flow_run.flow_id, name=flow_run.name, version=flow_run.version + 1, state=type(state).__name__, serialized_state=state.serialize(), ) flow = await models.Flow.where(id=flow_run.flow_id).first( selection_set={ "id": True, "tenant_id": True, "name": True, "environment": True, "version_group_id": True, }) tenant = await models.Tenant.where(id=flow_run.tenant_id ).first({"id", "slug"}) test_event = events.FlowRunStateChange( flow_run=new_flow_run.to_dict(), flow=flow.dict(), tenant=tenant.dict(), state=dict( version=1, state=serialized_state["type"], serialized_state=serialized_state, ), ) # otherwise populate the event with dummy data else: test_event = events.FlowRunStateChange( is_test_event=True, flow_run=dict(id=str(uuid.uuid4()), name=names.generate_slug(2)), tenant=dict(id=hook.tenant_id, slug="test-slug"), flow=dict(id=str(uuid.uuid4()), name="Cloud Hook Test Flow"), state=dict( version=1, state=serialized_state["type"], serialized_state=serialized_state, ), ) if hook.type == "WEBHOOK": await _call_webhook(url=hook.config["url"], event=test_event) elif hook.type == "SLACK_WEBHOOK": await _call_slack_webhook(url=hook.config["url"], event=test_event) elif hook.type == "PREFECT_MESSAGE": await _call_prefect_message(event=test_event) elif hook.type == "TWILIO": await _call_twilio( account_sid=hook.config["account_sid"], auth_token=hook.config["auth_token"], messaging_service_sid=hook.config["messaging_service_sid"], to=hook.config["to"], event=test_event, ) elif hook.type == "PAGERDUTY": await _call_pagerduty( api_token=hook.config["api_token"], routing_key=hook.config["routing_key"], severity=hook.config["severity"], event=test_event, ) else: raise ValueError("Invalid type")