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.")