def test_update_task_state(self): self._create_task_state("1", "2", self.execution_date) task_state = TaskState.get_task_state(self.dag_id, self.task_id, self.execution_date) task_state.task_state = 100 task_state.update_task_state() task_state = TaskState.get_task_state(self.dag_id, self.task_id, self.execution_date) assert task_state.task_state == 100
def test_get_task_state(self): dag_id = self.dag_id task_id = self.task_id task_state = TaskState.get_task_state(dag_id, task_id, self.execution_date) assert task_state.dag_id == dag_id assert task_state.task_id == task_id
def test_execute_event_handler(self): event = BaseEvent("test_event", "test_event", namespace="default") self.create_task_state(dag_run=self._dag_run, task_id='operator_toggle_handler') event_executor = DagRunEventExecutor(self._serialized_dag) actions = event_executor.execute_event_handler(self._dag_run, event=event) assert actions['operator_toggle_handler'] == SchedulingAction.START assert TaskState.get_task_state( dag_id="test_event_handler", task_id="operator_toggle_handler", executor_date=self._dag_run.execution_date).task_state is True actions = event_executor.execute_event_handler(self._dag_run, event=event) assert actions['operator_toggle_handler'] == SchedulingAction.STOP assert TaskState.get_task_state( dag_id="test_event_handler", task_id="operator_toggle_handler", executor_date=self._dag_run.execution_date).task_state is False
def _operator_handle_event(event, operator, execution_date) -> SchedulingAction: task_state = TaskState.get_task_state(operator.dag_id, operator.task_id, execution_date) event_handler = operator.get_events_handler() if task_state: scheduling_action, state = event_handler.handle_event( event, task_state.task_state) task_state.task_state = state task_state.update_task_state() else: scheduling_action, state = event_handler.handle_event(event, None) return scheduling_action
def _operator_handle_event(event, operator, execution_date) -> SchedulingAction: task_state = TaskState.get_task_state(operator.dag_id, operator.task_id, execution_date) event_handler = operator.get_events_handler() if not task_state: task_state = TaskState(task_id=operator.task_id, dag_id=operator.dag_id, execution_date=execution_date) with create_session() as session: session.add(task_state) session.commit() scheduling_action, state = event_handler.handle_event( event, task_state.task_state) task_state.task_state = state task_state.update_task_state() return scheduling_action