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.")
def start_new_workflow_execution( self, project_name: Text, workflow_name: Text) -> Optional[WorkflowExecutionInfo]: dag_id = self.airflow_dag_id(project_name, workflow_name) deploy_path = self.config.get('airflow_deploy_path') if deploy_path is None: raise Exception("airflow_deploy_path config not set!") if not self.dag_exist(dag_id): return None context: ExecutionContext = self.airflow_client.schedule_dag(dag_id) with create_session() as session: dagrun = DagRun.get_run_by_id(session=session, dag_id=dag_id, run_id=context.dagrun_id) if dagrun is None: return None else: return WorkflowExecutionInfo( workflow_info=WorkflowInfo(namespace=project_name, workflow_name=workflow_name), workflow_execution_id=str(dagrun.id), status=status.Status.INIT)