def test_set_dag_run_state_edge_cases(self): # Dag does not exist altered = set_dag_run_state_to_success(None, self.execution_dates[0]) self.assertEqual(len(altered), 0) altered = set_dag_run_state_to_failed(None, self.execution_dates[0]) self.assertEqual(len(altered), 0) altered = set_dag_run_state_to_running(None, self.execution_dates[0]) self.assertEqual(len(altered), 0) # Invalid execution date altered = set_dag_run_state_to_success(self.dag1, None) self.assertEqual(len(altered), 0) altered = set_dag_run_state_to_failed(self.dag1, None) self.assertEqual(len(altered), 0) altered = set_dag_run_state_to_running(self.dag1, None) self.assertEqual(len(altered), 0) # This will throw AssertionError since dag.latest_execution_date # need to be 0 does not exist. self.assertRaises(AssertionError, set_dag_run_state_to_success, self.dag1, timezone.make_naive(self.execution_dates[0])) # altered = set_dag_run_state_to_success(self.dag1, self.execution_dates[0]) # DagRun does not exist # This will throw AssertionError since dag.latest_execution_date does not exist self.assertRaises(AssertionError, set_dag_run_state_to_success, self.dag1, self.execution_dates[0])
def test_set_dag_run_state_edge_cases(self): # Dag does not exist altered = set_dag_run_state_to_success(None, self.execution_dates[0]) self.assertEqual(len(altered), 0) altered = set_dag_run_state_to_failed(None, self.execution_dates[0]) self.assertEqual(len(altered), 0) altered = set_dag_run_state_to_running(None, self.execution_dates[0]) self.assertEqual(len(altered), 0) # Invalid execution date altered = set_dag_run_state_to_success(self.dag1, None) self.assertEqual(len(altered), 0) altered = set_dag_run_state_to_failed(self.dag1, None) self.assertEqual(len(altered), 0) altered = set_dag_run_state_to_running(self.dag1, None) self.assertEqual(len(altered), 0) # This will throw AssertionError since dag.latest_execution_date # need to be 0 does not exist. self.assertRaises(AssertionError, set_dag_run_state_to_success, self.dag1, timezone.make_naive(self.execution_dates[0])) # altered = set_dag_run_state_to_success(self.dag1, self.execution_dates[0]) # DagRun does not exist # This will throw AssertionError since dag.latest_execution_date does not exist self.assertRaises(AssertionError, set_dag_run_state_to_success, self.dag1, self.execution_dates[0])
def test_set_state_without_commit(self): date = self.execution_dates[0] dr = self._create_test_dag_run(State.RUNNING, date) self._set_default_task_instance_states(dr) will_be_altered = set_dag_run_state_to_running(self.dag1, date, commit=False) # None of the tasks will be altered. self.assertEqual(len(will_be_altered), 0) self._verify_dag_run_state(self.dag1, date, State.RUNNING) self._verify_task_instance_states_remain_default(dr) will_be_altered = set_dag_run_state_to_failed(self.dag1, date, commit=False) # Only the running task will be altered. self.assertEqual(len(will_be_altered), 1) self._verify_dag_run_state(self.dag1, date, State.RUNNING) self._verify_task_instance_states_remain_default(dr) will_be_altered = set_dag_run_state_to_success(self.dag1, date, commit=False) # All except the SUCCESS task should be altered. self.assertEqual(len(will_be_altered), 5) self._verify_dag_run_state(self.dag1, date, State.RUNNING) self._verify_task_instance_states_remain_default(dr)
def test_set_state_without_commit(self): date = self.execution_dates[0] dr = self._create_test_dag_run(State.RUNNING, date) self._set_default_task_instance_states(dr) will_be_altered = set_dag_run_state_to_running(self.dag1, date, commit=False) # None of the tasks will be altered. self.assertEqual(len(will_be_altered), 0) self._verify_dag_run_state(self.dag1, date, State.RUNNING) self._verify_task_instance_states_remain_default(dr) will_be_altered = set_dag_run_state_to_failed(self.dag1, date, commit=False) # Only the running task will be altered. self.assertEqual(len(will_be_altered), 1) self._verify_dag_run_state(self.dag1, date, State.RUNNING) self._verify_task_instance_states_remain_default(dr) will_be_altered = set_dag_run_state_to_success(self.dag1, date, commit=False) # All except the SUCCESS task should be altered. self.assertEqual(len(will_be_altered), 5) self._verify_dag_run_state(self.dag1, date, State.RUNNING) self._verify_task_instance_states_remain_default(dr)
def test_set_state_without_commit(self): date = self.execution_dates[0] dr = self._create_test_dag_run(State.RUNNING, date) self._set_default_task_instance_states(dr) will_be_altered = set_dag_run_state_to_running(self.dag1, date, commit=False) # None of the tasks will be altered. assert len(will_be_altered) == 0 self._verify_dag_run_state(self.dag1, date, State.RUNNING) self._verify_task_instance_states_remain_default(dr) will_be_altered = set_dag_run_state_to_failed(self.dag1, date, commit=False) # Only the running task should be altered. expected = self._get_num_tasks_with_starting_state(State.RUNNING, inclusion=True) assert len(will_be_altered) == expected self._verify_dag_run_state(self.dag1, date, State.RUNNING) self._verify_task_instance_states_remain_default(dr) will_be_altered = set_dag_run_state_to_success(self.dag1, date, commit=False) # All except the SUCCESS task should be altered. expected = self._get_num_tasks_with_starting_state(State.SUCCESS, inclusion=False) assert len(will_be_altered) == expected self._verify_dag_run_state(self.dag1, date, State.RUNNING) self._verify_task_instance_states_remain_default(dr)
def update_dag_run_state(dag_id: str, dag_run_id: str, session) -> dict: """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 == State.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)
def test_set_failed_dag_run_to_success(self): date = self.execution_dates[0] dr = self._create_test_dag_run(State.SUCCESS, date) self._set_default_task_instance_states(dr) altered = set_dag_run_state_to_success(self.dag1, date, commit=True) # All except the SUCCESS task should be altered. self.assertEqual(len(altered), 5) self._verify_dag_run_state(self.dag1, date, State.SUCCESS) self._verify_task_instance_states(self.dag1, date, State.SUCCESS)
def test_set_dag_run_state_edge_cases(self): # Dag does not exist altered = set_dag_run_state_to_success(None, self.execution_dates[0]) assert len(altered) == 0 altered = set_dag_run_state_to_failed(None, self.execution_dates[0]) assert len(altered) == 0 altered = set_dag_run_state_to_running(None, self.execution_dates[0]) assert len(altered) == 0 # Invalid execution date altered = set_dag_run_state_to_success(self.dag1, None) assert len(altered) == 0 altered = set_dag_run_state_to_failed(self.dag1, None) assert len(altered) == 0 altered = set_dag_run_state_to_running(self.dag1, None) assert len(altered) == 0 # This will throw ValueError since dag.latest_execution_date # need to be 0 does not exist. with pytest.raises(ValueError): set_dag_run_state_to_success( self.dag2, timezone.make_naive(self.execution_dates[0])) # altered = set_dag_run_state_to_success(self.dag1, self.execution_dates[0]) # DagRun does not exist # This will throw ValueError since dag.latest_execution_date does not exist with pytest.raises(ValueError): set_dag_run_state_to_success(self.dag2, self.execution_dates[0])
def test_set_failed_dag_run_to_success(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_success(self.dag1, date, commit=True) # All except the SUCCESS task should be altered. self.assertEqual(len(altered), 5) self._verify_dag_run_state(self.dag1, date, State.SUCCESS) self._verify_task_instance_states(self.dag1, date, State.SUCCESS) self._verify_dag_run_dates(self.dag1, date, State.SUCCESS, middle_time)
def test_set_running_dag_run_to_success(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_success(self.dag1, date, commit=True) # All except the SUCCESS task should be altered. self.assertEqual(len(altered), 5) self._verify_dag_run_state(self.dag1, date, State.SUCCESS) self._verify_task_instance_states(self.dag1, date, State.SUCCESS) self._verify_dag_run_dates(self.dag1, date, State.SUCCESS, middle_time)
def test_set_failed_dag_run_to_success(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_success(self.dag1, date, commit=True) # All except the SUCCESS task should be altered. expected = self._get_num_tasks_with_starting_state(State.SUCCESS, inclusion=False) assert len(altered) == expected self._verify_dag_run_state(self.dag1, date, State.SUCCESS) self._verify_task_instance_states(self.dag1, date, State.SUCCESS) self._verify_dag_run_dates(self.dag1, date, State.SUCCESS, middle_time)
def test_set_state_with_multiple_dagruns(self, session=None): self.dag2.create_dagrun( run_type=DagRunType.MANUAL, state=State.FAILED, execution_date=self.execution_dates[0], session=session, ) self.dag2.create_dagrun( run_type=DagRunType.MANUAL, state=State.FAILED, execution_date=self.execution_dates[1], session=session, ) self.dag2.create_dagrun( run_type=DagRunType.MANUAL, state=State.RUNNING, execution_date=self.execution_dates[2], session=session, ) altered = set_dag_run_state_to_success(self.dag2, self.execution_dates[1], commit=True) # Recursively count number of tasks in the dag def count_dag_tasks(dag): count = len(dag.tasks) subdag_counts = [count_dag_tasks(subdag) for subdag in dag.subdags] count += sum(subdag_counts) return count self.assertEqual(len(altered), count_dag_tasks(self.dag2)) self._verify_dag_run_state(self.dag2, self.execution_dates[1], State.SUCCESS) # Make sure other dag status are not changed models.DagRun.find(dag_id=self.dag2.dag_id, execution_date=self.execution_dates[0]) self._verify_dag_run_state(self.dag2, self.execution_dates[0], State.FAILED) models.DagRun.find(dag_id=self.dag2.dag_id, execution_date=self.execution_dates[2]) self._verify_dag_run_state(self.dag2, self.execution_dates[2], State.RUNNING)
def test_set_state_with_multiple_dagruns(self): dr1 = self.dag2.create_dagrun( run_id='manual__' + datetime.now().isoformat(), state=State.FAILED, execution_date=self.execution_dates[0], session=self.session ) dr2 = self.dag2.create_dagrun( run_id='manual__' + datetime.now().isoformat(), state=State.FAILED, execution_date=self.execution_dates[1], session=self.session ) dr3 = self.dag2.create_dagrun( run_id='manual__' + datetime.now().isoformat(), state=State.RUNNING, execution_date=self.execution_dates[2], session=self.session ) altered = set_dag_run_state_to_success(self.dag2, self.execution_dates[1], commit=True) # Recursively count number of tasks in the dag def count_dag_tasks(dag): count = len(dag.tasks) subdag_counts = [count_dag_tasks(subdag) for subdag in dag.subdags] count += sum(subdag_counts) return count self.assertEqual(len(altered), count_dag_tasks(self.dag2)) self._verify_dag_run_state(self.dag2, self.execution_dates[1], State.SUCCESS) # Make sure other dag status are not changed dr1 = models.DagRun.find(dag_id=self.dag2.dag_id, execution_date=self.execution_dates[0]) dr1 = dr1[0] self._verify_dag_run_state(self.dag2, self.execution_dates[0], State.FAILED) dr3 = models.DagRun.find(dag_id=self.dag2.dag_id, execution_date=self.execution_dates[2]) dr3 = dr3[0] self._verify_dag_run_state(self.dag2, self.execution_dates[2], State.RUNNING)