def update_states(self, flow_runs: list) -> None:
"""
After a flow run is grabbed this function sets the state to Submitted so it
won't be picked up by any other processes
Args:
- flow_runs (list): A list of GraphQLResult flow run objects
"""
for flow_run in flow_runs:
# Set flow run state to `Submitted` if it is currently `Scheduled`
if state.StateSchema().load(flow_run.serialized_state).is_scheduled():
self.client.set_flow_run_state(
flow_run_id=flow_run.id,
version=flow_run.version,
state=Submitted(
message="Submitted for execution",
state=state.StateSchema().load(flow_run.serialized_state),
),
)
# Set task run states to `Submitted` if they are currently `Scheduled`
for task_run in flow_run.task_runs:
if state.StateSchema().load(task_run.serialized_state).is_scheduled():
self.client.set_task_run_state(
task_run_id=task_run.id,
version=task_run.version,
state=Submitted(
message="Submitted for execution",
state=state.StateSchema().load(task_run.serialized_state),
),
)
def update_state(self, flow_run: GraphQLResult, deployment_info: str) -> None:
"""
After a flow run is grabbed this function sets the state to Submitted so it
won't be picked up by any other processes
Args:
- flow_run (GraphQLResult): A GraphQLResult flow run object
- deployment_info (str): Identifier information related to the Flow Run
deployment
"""
self.logger.debug(
"Updating states for flow run {}".format(flow_run.id) # type: ignore
)
# Set flow run state to `Submitted` if it is currently `Scheduled`
if state.StateSchema().load(flow_run.serialized_state).is_scheduled():
self.logger.debug(
"Flow run {} is in a Scheduled state, updating to Submitted".format(
flow_run.id # type: ignore
)
)
self.client.set_flow_run_state(
flow_run_id=flow_run.id,
version=flow_run.version,
state=Submitted(
message="Submitted for execution. {}".format(deployment_info),
state=state.StateSchema().load(flow_run.serialized_state),
),
)
# Set task run states to `Submitted` if they are currently `Scheduled`
for task_run in flow_run.task_runs:
if state.StateSchema().load(task_run.serialized_state).is_scheduled():
self.logger.debug(
"Task run {} is in a Scheduled state, updating to Submitted".format(
task_run.id # type: ignore
)
)
self.client.set_task_run_state(
task_run_id=task_run.id,
version=task_run.version,
state=Submitted(
message="Submitted for execution. {}".format(deployment_info),
state=state.StateSchema().load(task_run.serialized_state),
),
)
def _mark_flow_as_submitted(self, flow_run: GraphQLResult) -> None:
"""
After a flow run is grabbed this function sets the state to Submitted so it
won't be picked up by any other processes
Args:
- flow_run (GraphQLResult): A GraphQLResult flow run object
"""
# Set flow run state to `Submitted` if it is currently `Scheduled`
if state.StateSchema().load(flow_run.serialized_state).is_scheduled():
self.logger.debug(
f"Updating flow run {flow_run.id} state from Scheduled -> Submitted..."
)
self.client.set_flow_run_state(
flow_run_id=flow_run.id,
version=flow_run.version,
state=Submitted(
message="Submitted for execution",
state=state.StateSchema().load(flow_run.serialized_state),
),
)
# Set task run states to `Submitted` if they are currently `Scheduled`
task_runs_updated = 0
for task_run in flow_run.task_runs:
if state.StateSchema().load(
task_run.serialized_state).is_scheduled():
task_runs_updated += 1
self.client.set_task_run_state(
task_run_id=task_run.id,
version=task_run.version,
state=Submitted(
message="Submitted for execution.",
state=state.StateSchema().load(
task_run.serialized_state),
),
)
if task_runs_updated:
self.logger.debug(
f"Updated {task_runs_updated} task runs states for flow run "
f"{flow_run.id} from Scheduled -> Submitted")
