def _send_scheduling_task_event(self, ti: Optional[TI], action: SchedulingAction): if ti is None: return task_scheduling_event = TaskSchedulingEvent(ti.task_id, ti.dag_id, ti.execution_date, ti.try_number, action) self.mailbox.send_message(task_scheduling_event.to_event())
def test_scheduling_task(self): e_date = timezone.utcnow() scheduling_task = TaskSchedulingEvent(task_id='a', dag_id='b', execution_date=e_date, try_number=1, action=SchedulingAction.START) event = TaskSchedulingEvent.to_base_event(scheduling_task) scheduling_task_2 = TaskSchedulingEvent.from_base_event(event) self.assertEqual(e_date, scheduling_task_2.execution_date) self.assertEqual(SchedulingAction.START, scheduling_task_2.action)
def _process_request_event(self, event: RequestEvent, session: Session = None): try: message = BaseUserDefineMessage() message.from_json(event.body) if message.message_type == UserDefineMessageType.RUN_DAG: # todo make sure dag file is parsed. dagrun = self._create_dag_run(message.dag_id, session=session, run_type=DagRunType.MANUAL) if not dagrun: self.log.error("Failed to create dag_run.") # TODO Need to add ret_code and errro_msg in ExecutionContext in case of exception self.notification_client.send_event( ResponseEvent(event.request_id, None).to_event()) return tasks = self._find_schedulable_tasks(dagrun, session, False) self._send_scheduling_task_events(tasks, SchedulingAction.START) self.notification_client.send_event( ResponseEvent(event.request_id, dagrun.run_id).to_event()) elif message.message_type == UserDefineMessageType.STOP_DAG_RUN: dag_run = DagRun.get_run_by_id(session=session, dag_id=message.dag_id, run_id=message.dagrun_id) self._stop_dag_run(dag_run) self.notification_client.send_event( ResponseEvent(event.request_id, dag_run.run_id).to_event()) elif message.message_type == UserDefineMessageType.EXECUTE_TASK: dagrun = DagRun.get_run_by_id(session=session, dag_id=message.dag_id, run_id=message.dagrun_id) ti: TI = dagrun.get_task_instance(task_id=message.task_id) self.mailbox.send_message( TaskSchedulingEvent(task_id=ti.task_id, dag_id=ti.dag_id, execution_date=ti.execution_date, try_number=ti.try_number, action=SchedulingAction( message.action)).to_event()) self.notification_client.send_event( ResponseEvent(event.request_id, dagrun.run_id).to_event()) except Exception: self.log.exception("Error occurred when processing request event.")