async def test_nested_insert_array(self, flow_id):
""" insert nested objects as an array"""
flow_run_id = await models.FlowRun(
flow_id=flow_id,
states=[
models.FlowRunState(state="test", serialized_state={}),
models.FlowRunState(state="test", serialized_state={}),
],
).insert()
assert (await models.FlowRunState.where({
"flow_run_id": {
"_eq": flow_run_id
}
}).count() == 2)
async def _create_flow_run(
flow_id: str = None,
parameters: dict = None,
context: dict = None,
scheduled_start_time: datetime.datetime = None,
flow_run_name: str = None,
version_group_id: str = None,
labels: List[str] = None,
run_config: dict = None,
) -> Any:
"""
Creates a new flow run for an existing flow.
Args:
- flow_id (str): A string representing the current flow id
- parameters (dict, optional): A dictionary of parameters that were specified for the flow
- context (dict, optional): A dictionary of context values
- scheduled_start_time (datetime.datetime): When the flow_run should be scheduled to run. If `None`,
defaults to right now. Must be UTC.
- flow_run_name (str, optional): An optional string representing this flow run
- version_group_id (str, optional): An optional version group ID; if provided, will run the most
recent unarchived version of the group
- labels (List[str], optional): a list of labels to apply to this individual flow run
- run-config (dict, optional): A run-config override for this flow run.
"""
if flow_id is None and version_group_id is None:
raise ValueError(
"One of flow_id or version_group_id must be provided.")
scheduled_start_time = scheduled_start_time or pendulum.now()
if flow_id:
where_clause = {"id": {"_eq": flow_id}}
elif version_group_id:
where_clause = {
"version_group_id": {
"_eq": version_group_id
},
"archived": {
"_eq": False
},
}
flow = await models.Flow.where(where=where_clause).first(
{
"id": True,
"archived": True,
"tenant_id": True,
"environment": True,
"run_config": True,
"parameters": True,
"flow_group_id": True,
"flow_group": {
"default_parameters": True,
"labels": True,
"run_config": True,
},
},
order_by={"version": EnumValue("desc")},
) # type: Any
if not flow:
msg = (f"Flow {flow_id} not found" if flow_id else
f"Version group {version_group_id} has no unarchived flows.")
raise exceptions.NotFound(msg)
elif flow.archived:
raise ValueError(f"Flow {flow.id} is archived.")
# determine active labels
if labels is not None:
run_labels = labels
elif run_config is not None:
run_labels = run_config.get("labels") or []
elif flow.flow_group.labels is not None:
run_labels = flow.flow_group.labels
elif flow.flow_group.run_config is not None:
run_labels = flow.flow_group.run_config.get("labels") or []
elif flow.run_config is not None:
run_labels = flow.run_config.get("labels") or []
elif flow.environment is not None:
run_labels = flow.environment.get("labels") or []
else:
run_labels = []
run_labels.sort()
# determine active run_config
if run_config is None:
if flow.flow_group.run_config is not None:
run_config = flow.flow_group.run_config
else:
run_config = flow.run_config
# check parameters
run_parameters = flow.flow_group.default_parameters
run_parameters.update((parameters or {}))
required_parameters = [p["name"] for p in flow.parameters if p["required"]]
missing = set(required_parameters).difference(run_parameters)
if missing:
raise ValueError(f"Required parameters were not supplied: {missing}")
state = Scheduled(message="Flow run scheduled.",
start_time=scheduled_start_time)
run = models.FlowRun(
tenant_id=flow.tenant_id,
flow_id=flow_id or flow.id,
labels=run_labels,
parameters=run_parameters,
run_config=run_config,
context=context or {},
scheduled_start_time=scheduled_start_time,
name=flow_run_name or names.generate_slug(2),
states=[
models.FlowRunState(
tenant_id=flow.tenant_id,
**models.FlowRunState.fields_from_state(
Pending(message="Flow run created")),
)
],
)
flow_run_id = await run.insert()
# apply the flow run's initial state via `set_flow_run_state`
await api.states.set_flow_run_state(flow_run_id=flow_run_id, state=state)
return flow_run_id
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