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")