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)