def list_workflow_executions(
self, project_name: Text,
workflow_name: Text) -> List[WorkflowExecutionInfo]:
return [
WorkflowExecutionInfo(workflow_execution_id='1',
status=Status.RUNNING),
WorkflowExecutionInfo(workflow_execution_id='2',
status=Status.RUNNING)
]
def test_list_workflow_executions(self):
with mock.patch(SCHEDULER_CLASS) as mockScheduler:
instance = mockScheduler.return_value
self.server.scheduler_service._scheduler = instance
instance.list_workflow_executions.return_value \
= [WorkflowExecutionInfo(workflow_execution_id='id_1', status=Status.INIT),
WorkflowExecutionInfo(workflow_execution_id='id_2', status=Status.INIT)]
client = SchedulerClient("localhost:{}".format(_PORT))
workflow_execution_list = client.list_workflow_executions(
namespace='namespace', workflow_name='test_workflow')
self.assertEqual(2, len(workflow_execution_list))
def restart_job_execution(
self, job_name: Text,
workflow_execution_id: Text) -> Optional[JobExecutionInfo]:
with create_session() as session:
dagrun = session.query(DagRun).filter(
DagRun.id == int(workflow_execution_id)).first()
if dagrun is None:
return None
if dagrun.state != State.RUNNING:
raise Exception(
'execution: {} state: {} can not trigger job.'.format(
workflow_execution_id, dagrun.state))
task = dagrun.get_task_instance(job_name, session)
if task is None:
return None
self.airflow_client.schedule_task(
dag_id=dagrun.dag_id,
task_id=job_name,
action=SchedulingAction.RESTART,
context=ExecutionContext(dagrun_id=dagrun.run_id))
project_name, workflow_name = self.dag_id_to_namespace_workflow(
dagrun.dag_id)
return JobExecutionInfo(
job_name=job_name,
status=self.airflow_state_to_status(task.state),
workflow_execution=WorkflowExecutionInfo(
workflow_info=WorkflowInfo(namespace=project_name,
workflow_name=workflow_name),
workflow_execution_id=workflow_execution_id,
status=self.airflow_state_to_status(dagrun.state)))
def stop_job_execution(
self, job_name: Text,
workflow_execution_id: Text) -> Optional[JobExecutionInfo]:
with create_session() as session:
dagrun = session.query(DagRun).filter(
DagRun.id == int(workflow_execution_id)).first()
if dagrun is None:
return None
task = dagrun.get_task_instance(job_name, session)
if task is None:
return None
if task.state in State.finished:
raise Exception('job:{} state: {} can not stop!'.format(
job_name, task.state))
else:
self.airflow_client.schedule_task(
dag_id=dagrun.dag_id,
task_id=job_name,
action=SchedulingAction.STOP,
context=ExecutionContext(dagrun_id=dagrun.run_id))
project_name, workflow_name = self.dag_id_to_namespace_workflow(
dagrun.dag_id)
return JobExecutionInfo(
job_name=job_name,
status=self.airflow_state_to_status(task.state),
workflow_execution=WorkflowExecutionInfo(
workflow_info=WorkflowInfo(namespace=project_name,
workflow_name=workflow_name),
workflow_execution_id=workflow_execution_id,
status=self.airflow_state_to_status(dagrun.state)))
def list_workflow_executions(
self, project_name: Text,
workflow_name: Text) -> List[WorkflowExecutionInfo]:
dag_id = self.airflow_dag_id(project_name, workflow_name)
with create_session() as session:
dagrun_list = session.query(DagRun).filter(
DagRun.dag_id == dag_id).all()
if dagrun_list is None:
return []
else:
result = []
for dagrun in dagrun_list:
status_ = self.airflow_state_to_status(dagrun.state)
result.append(
WorkflowExecutionInfo(
workflow_info=WorkflowInfo(
namespace=project_name,
workflow_name=workflow_name),
workflow_execution_id=str(dagrun.id),
status=status_,
start_date=str(datetime_to_int64(
dagrun.start_date)),
end_date=str(datetime_to_int64(dagrun.end_date)),
))
return result
def start_job_execution(self, job_name: Text,
execution_id: Text) -> JobExecutionInfo:
job = self.workflow.jobs[job_name]
plugins = self.workflow.properties.get(
WorkflowPropertyKeys.JOB_PLUGINS)
module, name = plugins.get(job.job_config.job_type)
class_object = import_string('{}.{}'.format(module, name))
self.job_controller: JobController = class_object()
job_execution_info = JobExecutionInfo(
job_execution_id='1',
job_name=job_name,
workflow_execution=WorkflowExecutionInfo(
workflow_execution_id='1',
workflow_info=WorkflowInfo(
workflow_name=self.workflow.workflow_name)))
self.job_runtime_env: JobRuntimeEnv = prepare_job_runtime_env(
root_working_dir=self.project_context.project_path + '/temp',
workflow_snapshot_id=self.workflow.workflow_snapshot_id,
workflow_name=self.workflow.workflow_name,
job_execution_info=job_execution_info,
project_context=self.project_context)
self.job_handler = self.job_controller.submit_job(
job=job, job_runtime_env=self.job_runtime_env)
return job_execution_info
def test_list_jobs(self):
with mock.patch(SCHEDULER_CLASS) as mockScheduler:
instance = mockScheduler.return_value
self.server.scheduler_service._scheduler = instance
instance.list_job_executions.return_value \
= [JobExecutionInfo(job_name='job_name_1',
status=Status.RUNNING,
workflow_execution=WorkflowExecutionInfo(workflow_execution_id='id',
status=Status.INIT)),
JobExecutionInfo(job_name='job_name_2',
status=Status.RUNNING,
workflow_execution=WorkflowExecutionInfo(workflow_execution_id='id',
status=Status.INIT))]
client = SchedulerClient("localhost:{}".format(_PORT))
job_list = client.list_jobs(execution_id='id')
self.assertEqual(2, len(job_list))
def stop_workflow_execution(
self, execution_id: Text) -> Optional[WorkflowExecutionInfo]:
workflow_info = WorkflowInfo(workflow_name='workflow_name',
namespace='project_name')
workflow_execution_info = WorkflowExecutionInfo(
workflow_execution_id='1',
workflow_info=workflow_info,
status=Status.KILLED,
start_date=str(int(time.time() * 1000)))
return workflow_execution_info
def proto_to_workflow_execution(proto: WorkflowExecutionProto) -> WorkflowExecutionInfo:
if proto is None:
return None
else:
return WorkflowExecutionInfo(workflow_execution_id=proto.execution_id,
status=proto_to_state(proto.execution_state),
start_date=str(proto.start_time.value),
end_date=str(proto.end_time.value),
workflow_info=proto_to_workflow(proto.workflow),
properties=dict(proto.properties))
def stop_all_workflow_execution(
self, project_name: Text,
workflow_name: Text) -> List[WorkflowExecutionInfo]:
workflow_info = WorkflowInfo(workflow_name=workflow_name,
namespace=project_name)
workflow_execution_info = WorkflowExecutionInfo(
workflow_execution_id='1',
workflow_info=workflow_info,
status=Status.KILLED,
start_date=str(int(time.time() * 1000)))
return [workflow_execution_info]
def start_new_workflow_execution(
self, project_name: Text,
workflow_name: Text) -> Optional[WorkflowExecutionInfo]:
workflow_info = WorkflowInfo(workflow_name=workflow_name,
namespace=project_name)
workflow_execution_info = WorkflowExecutionInfo(
workflow_execution_id='1',
workflow_info=workflow_info,
status=Status.RUNNING,
start_date=str(int(time.time() * 1000)))
return workflow_execution_info
def test_get_workflow_execution(self):
with mock.patch(SCHEDULER_CLASS) as mockScheduler:
instance = mockScheduler.return_value
self.server.scheduler_service._scheduler = instance
instance.get_workflow_execution.return_value \
= WorkflowExecutionInfo(workflow_execution_id='id', status=Status.INIT)
client = SchedulerClient("localhost:{}".format(_PORT))
workflow_execution = client.get_workflow_execution(
execution_id='id')
self.assertEqual('id', workflow_execution.execution_id)
self.assertEqual(StateProto.INIT,
workflow_execution.execution_state)
def list_job_executions(self,
execution_id: Text) -> List[JobExecutionInfo]:
workflow_info = WorkflowInfo(workflow_name='workflow_name',
namespace='project_name')
workflow_execution_info = WorkflowExecutionInfo(
workflow_execution_id='1',
workflow_info=workflow_info,
status=Status.KILLED,
start_date=str(int(time.time() * 1000)))
job_execution_info: JobExecutionInfo = JobExecutionInfo(
job_name='job_name',
status=Status.RUNNING,
workflow_execution=workflow_execution_info,
start_date=str(int(time.time() * 1000)))
return [job_execution_info]
def stop_job_execution(self, job_name: Text,
execution_id: Text) -> JobExecutionInfo:
workflow_info = WorkflowInfo(workflow_name='workflow_name',
namespace='project_name')
workflow_execution_info = WorkflowExecutionInfo(
workflow_execution_id='1',
workflow_info=workflow_info,
status=Status.KILLED,
start_date=str(int(time.time() * 1000)))
job_execution_info: JobExecutionInfo = JobExecutionInfo(
job_name=job_name,
status=Status.KILLED,
workflow_execution=workflow_execution_info,
start_date=str(int(time.time() * 1000)))
return job_execution_info
def test_init_job_runtime_context(self):
working_dir = os.path.dirname(__file__)
job_runtime_env = JobRuntimeEnv(
working_dir=working_dir,
job_execution_info=JobExecutionInfo(
job_name='task_1',
workflow_execution=WorkflowExecutionInfo(
workflow_execution_id='1',
workflow_info=WorkflowInfo(workflow_name='workflow_1'))))
init_job_runtime_context(job_runtime_env)
self.assertEqual('workflow_1', current_workflow_config().workflow_name)
self.assertEqual(
'task_1',
current_workflow_config().job_configs[current_job_name()].job_name)
self.assertEqual('test_project',
current_project_config().get_project_name())
def test_get_job(self):
with mock.patch(SCHEDULER_CLASS) as mockScheduler:
instance = mockScheduler.return_value
self.server.scheduler_service._scheduler = instance
instance.get_job_executions.return_value \
= [JobExecutionInfo(job_name='job_name',
status=Status.RUNNING,
workflow_execution=WorkflowExecutionInfo(workflow_execution_id='id',
status=Status.INIT))]
client = SchedulerClient("localhost:{}".format(_PORT))
job = client.get_job(job_name='job_name', execution_id='id')
self.assertEqual('job_name', job.name)
self.assertEqual(StateProto.RUNNING, job.job_state)
self.assertEqual('id', job.workflow_execution.execution_id)
self.assertEqual(StateProto.INIT,
job.workflow_execution.execution_state)
def build_job_execution_info_list(self, dagrun, task_list):
project_name, workflow_name = self.dag_id_to_namespace_workflow(
dagrun.dag_id)
result = []
for task in task_list:
job = JobExecutionInfo(
job_name=task.task_id,
status=self.airflow_state_to_status(task.state),
start_date=str(datetime_to_int64(task.start_date)),
end_date=str(datetime_to_int64(task.end_date)),
workflow_execution=WorkflowExecutionInfo(
workflow_info=WorkflowInfo(namespace=project_name,
workflow_name=workflow_name),
workflow_execution_id=str(dagrun.id),
status=self.airflow_state_to_status(dagrun.state)))
result.append(job)
return result
def stop_workflow_execution(
self,
workflow_execution_id: Text) -> Optional[WorkflowExecutionInfo]:
with create_session() as session:
dagrun = session.query(DagRun).filter(
DagRun.id == int(workflow_execution_id)).first()
if dagrun is None:
return None
project_name, workflow_name = self.dag_id_to_namespace_workflow(
dagrun.dag_id)
context: ExecutionContext = ExecutionContext(
dagrun_id=dagrun.run_id)
current_context = self.airflow_client.stop_dag_run(
dagrun.dag_id, context)
return WorkflowExecutionInfo(
workflow_info=WorkflowInfo(namespace=project_name,
workflow_name=workflow_name),
workflow_execution_id=workflow_execution_id,
status=status.Status.KILLED)
def get_workflow_execution(
self,
workflow_execution_id: Text) -> Optional[WorkflowExecutionInfo]:
with create_session() as session:
dagrun = session.query(DagRun).filter(
DagRun.id == int(workflow_execution_id)).first()
if dagrun is None:
return None
else:
status_ = self.airflow_state_to_status(dagrun.state)
project_name, workflow_name = self.dag_id_to_namespace_workflow(
dagrun.dag_id)
return WorkflowExecutionInfo(
workflow_info=WorkflowInfo(namespace=project_name,
workflow_name=workflow_name),
workflow_execution_id=workflow_execution_id,
status=status_,
start_date=str(datetime_to_int64(dagrun.start_date)),
end_date=str(datetime_to_int64(dagrun.end_date)))
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)
def context_to_job_info(self, project_name: Text,
context: Any) -> JobExecutionInfo:
"""
The function of this function is to turn the context of airflow into execution information of a job.
"""
wi = WorkflowInfo(namespace=project_name,
workflow_name=self.workflow.workflow_name)
we = WorkflowExecutionInfo(
workflow_execution_id=str(context.get('dag_run').id),
workflow_info=wi,
start_date=str(datetime_to_int64(
context.get('dag_run').start_date)),
end_date=str(datetime_to_int64(context.get('dag_run').end_date)),
status=context.get('dag_run').get_state())
je = JobExecutionInfo(
job_name=self.job.job_name,
job_execution_id=str(context.get('ti').try_number),
status=context.get('ti').state,
start_date=str(datetime_to_int64(context.get('ti').start_date)),
end_date=str(datetime_to_int64(context.get('ti').end_date)),
workflow_execution=we)
return je
def get_workflow_execution(
self,
workflow_execution_id: Text) -> Optional[WorkflowExecutionInfo]:
return WorkflowExecutionInfo(workflow_execution_id='1',
status=Status.RUNNING)
def start_new_workflow_execution(
self, project_name: Text,
workflow_name: Text) -> Optional[WorkflowExecutionInfo]:
return WorkflowExecutionInfo(workflow_execution_id='1',
status=Status.RUNNING)
