def test_set_success_dag_run_to_failed(self):
date = self.execution_dates[0]
dr = self._create_test_dag_run(State.SUCCESS, date)
middle_time = timezone.utcnow()
self._set_default_task_instance_states(dr)
altered = set_dag_run_state_to_failed(self.dag1, date, commit=True)
# Only running task should be altered.
expected = self._get_num_tasks_with_starting_state(State.RUNNING,
inclusion=True)
assert len(altered) == expected
self._verify_dag_run_state(self.dag1, date, State.FAILED)
assert dr.get_task_instance('run_after_loop').state == State.FAILED
self._verify_dag_run_dates(self.dag1, date, State.FAILED, middle_time)
def test_set_running_dag_run_to_failed(self):
date = self.execution_dates[0]
dr = self._create_test_dag_run(State.RUNNING, date)
middle_time = timezone.utcnow()
self._set_default_task_instance_states(dr)
altered = set_dag_run_state_to_failed(self.dag1, date, commit=True)
# Only running task should be altered.
self.assertEqual(len(altered), 1)
self._verify_dag_run_state(self.dag1, date, State.FAILED)
self.assertEqual(
dr.get_task_instance('run_after_loop').state, State.FAILED)
self._verify_dag_run_dates(self.dag1, date, State.FAILED, middle_time)
def update_dag_run_state(*,
dag_id: str,
dag_run_id: str,
session: Session = NEW_SESSION) -> APIResponse:
"""Set a state of a dag run."""
dag_run: Optional[DagRun] = (session.query(DagRun).filter(
DagRun.dag_id == dag_id, DagRun.run_id == dag_run_id).one_or_none())
if dag_run is None:
error_message = f'Dag Run id {dag_run_id} not found in dag {dag_id}'
raise NotFound(error_message)
try:
post_body = set_dagrun_state_form_schema.load(request.json)
except ValidationError as err:
raise BadRequest(detail=str(err))
state = post_body['state']
dag = current_app.dag_bag.get_dag(dag_id)
if state == DagRunState.SUCCESS:
set_dag_run_state_to_success(dag, dag_run.execution_date, commit=True)
else:
set_dag_run_state_to_failed(dag, dag_run.execution_date, commit=True)
dag_run = session.query(DagRun).get(dag_run.id)
return dagrun_schema.dump(dag_run)
