def test_should_render_dag_orientation(self): orientation = "TB" dag = DAG(dag_id="DAG_ID", orientation=orientation) task_1 = BashOperator(dag=dag, start_date=START_DATE, task_id="first", bash_command="echo 1") task_2 = BashOperator(dag=dag, start_date=START_DATE, task_id="second", bash_command="echo 1") task_3 = PythonOperator( dag=dag, start_date=START_DATE, task_id="third", python_callable=mock.MagicMock() ) task_1 >> task_2 task_1 >> task_3 tis = [ TaskInstance(task_1, execution_date=START_DATE, state=State.SCHEDULED), TaskInstance(task_2, execution_date=START_DATE, state=State.SUCCESS), TaskInstance(task_3, execution_date=START_DATE, state=State.RUNNING), ] dot = dot_renderer.render_dag(dag, tis=tis) source = dot.source # Should render DAG title with orientation assert "label=DAG_ID" in source assert f'label=DAG_ID labelloc=t rankdir={orientation}' in source # Change orientation orientation = "LR" dag = DAG(dag_id="DAG_ID", orientation=orientation) dot = dot_renderer.render_dag(dag, tis=tis) source = dot.source # Should render DAG title with orientation assert "label=DAG_ID" in source assert f'label=DAG_ID labelloc=t rankdir={orientation}' in source
def dag_show(args): """Displays DAG or saves it's graphic representation to the file""" dag = get_dag(args.subdir, args.dag_id) dot = render_dag(dag) if args.save: filename, _, fileformat = args.save.rpartition('.') dot.render(filename=filename, format=fileformat, cleanup=True) print("File {} saved".format(args.save)) elif args.imgcat: data = dot.pipe(format='png') try: proc = subprocess.Popen("imgcat", stdout=subprocess.PIPE, stdin=subprocess.PIPE) except OSError as e: if e.errno == errno.ENOENT: raise AirflowException( "Failed to execute. Make sure the imgcat executables are on your systems \'PATH\'" ) else: raise out, err = proc.communicate(data) if out: print(out.decode('utf-8')) if err: print(err.decode('utf-8')) else: print(dot.source)
def dag_test(args, session=None): """Execute one single DagRun for a given DAG and execution date, using the DebugExecutor.""" dag = get_dag(subdir=args.subdir, dag_id=args.dag_id) dag.clear(start_date=args.execution_date, end_date=args.execution_date, dag_run_state=State.NONE) try: dag.run(executor=DebugExecutor(), start_date=args.execution_date, end_date=args.execution_date) except BackfillUnfinished as e: print(str(e)) show_dagrun = args.show_dagrun imgcat = args.imgcat_dagrun filename = args.save_dagrun if show_dagrun or imgcat or filename: tis = session.query(TaskInstance).filter( TaskInstance.dag_id == args.dag_id, TaskInstance.execution_date == args.execution_date, ).all() dot_graph = render_dag(dag, tis=tis) print() if filename: _save_dot_to_file(dot_graph, filename) if imgcat: _display_dot_via_imgcat(dot_graph) if show_dagrun: print(dot_graph.source)
def test_should_render_dag(self): dag = DAG(dag_id="DAG_ID") task_1 = BashOperator(dag=dag, start_date=START_DATE, task_id="first", bash_command="echo 1") task_2 = BashOperator(dag=dag, start_date=START_DATE, task_id="second", bash_command="echo 1") task_3 = PythonOperator(dag=dag, start_date=START_DATE, task_id="third", python_callable=mock.MagicMock()) task_1 >> task_2 task_1 >> task_3 dot = dot_renderer.render_dag(dag) source = dot.source # Should render DAG title self.assertIn("label=DAG_ID", source) self.assertIn("first", source) self.assertIn("second", source) self.assertIn("third", source) self.assertIn("first -> second", source) self.assertIn("first -> third", source) self.assertIn('fillcolor="#f0ede4"', source) self.assertIn('fillcolor="#f0ede4"', source)
def test_should_render_dag_with_task_instances(self): dag = DAG(dag_id="DAG_ID") task_1 = BashOperator(dag=dag, start_date=START_DATE, task_id="first", bash_command="echo 1") task_2 = BashOperator(dag=dag, start_date=START_DATE, task_id="second", bash_command="echo 1") task_3 = PythonOperator( dag=dag, start_date=START_DATE, task_id="third", python_callable=mock.MagicMock() ) task_1 >> task_2 task_1 >> task_3 tis = [ TaskInstance(task_1, execution_date=START_DATE, state=State.SCHEDULED), TaskInstance(task_2, execution_date=START_DATE, state=State.SUCCESS), TaskInstance(task_3, execution_date=START_DATE, state=State.RUNNING), ] dot = dot_renderer.render_dag(dag, tis=tis) source = dot.source # Should render DAG title assert "label=DAG_ID" in source assert ( 'first [color=black fillcolor=tan label=first shape=rectangle style="filled,rounded"]' in source ) assert ( 'second [color=white fillcolor=green label=second shape=rectangle style="filled,rounded"]' in source ) assert ( 'third [color=black fillcolor=lime label=third shape=rectangle style="filled,rounded"]' in source )
def test_disable_task(self): dag_creator = DagCreator(self.task_graph._graph, with_data_nodes=True) dags = dag_creator.traverse_graph() self.assertEqual(len(dags), 3) test_spark_dag = dags['test_spark'] dot = render_dag(test_spark_dag) self.assertEqual(dot.source, self.dot_test_spark_deactivate)
def test_dag_creator_external_sensor(self): dag_creator = DagCreator(self.task_graph._graph) dags = dag_creator.traverse_graph() self.assertEqual(len(dags), 3) test_external_sensor_dag = dags['test_external_sensor'] dot = render_dag(test_external_sensor_dag) self.assertEqual(dot.source, self.dot_test_external_sensor)
def test_dag_creator_with_dataset(self): dag_creator = DagCreator(self.task_graph._graph, with_data_nodes=True) dags = dag_creator.traverse_graph() self.assertEqual(len(dags), 3) test_batch_dag = dags['test_batch'] dot = render_dag(test_batch_dag) self.assertEqual(dot.source, self.dot_test_batch_graph_with_dataset)
def dag_show(args): """Displays DAG or saves it's graphic representation to the file""" dag = get_dag(args.subdir, args.dag_id) dot = render_dag(dag) filename = args.save imgcat = args.imgcat if filename and imgcat: raise SystemExit( "Option --save and --imgcat are mutually exclusive. " "Please remove one option to execute the command.", ) elif filename: _save_dot_to_file(dot, filename) elif imgcat: _display_dot_via_imgcat(dot) else: print(dot.source)
def dag_test(args, session=None): """Execute one single DagRun for a given DAG and execution date, using the DebugExecutor.""" dag = get_dag(subdir=args.subdir, dag_id=args.dag_id) dag.clear(start_date=args.execution_date, end_date=args.execution_date, dag_run_state=False) try: dag.run( executor=DebugExecutor(), start_date=args.execution_date, end_date=args.execution_date, # Always run the DAG at least once even if no logical runs are # available. This does not make a lot of sense, but Airflow has # been doing this prior to 2.2 so we keep compatibility. run_at_least_once=True, ) except BackfillUnfinished as e: print(str(e)) show_dagrun = args.show_dagrun imgcat = args.imgcat_dagrun filename = args.save_dagrun if show_dagrun or imgcat or filename: tis = ( session.query(TaskInstance) .filter( TaskInstance.dag_id == args.dag_id, TaskInstance.execution_date == args.execution_date, ) .all() ) dot_graph = render_dag(dag, tis=tis) print() if filename: _save_dot_to_file(dot_graph, filename) if imgcat: _display_dot_via_imgcat(dot_graph) if show_dagrun: print(dot_graph.source)
def test_render_task_group(self): with DAG(dag_id="example_task_group", start_date=START_DATE) as dag: start = DummyOperator(task_id="start") with TaskGroup("section_1", tooltip="Tasks for section_1") as section_1: task_1 = DummyOperator(task_id="task_1") task_2 = BashOperator(task_id="task_2", bash_command='echo 1') task_3 = DummyOperator(task_id="task_3") task_1 >> [task_2, task_3] with TaskGroup("section_2", tooltip="Tasks for section_2") as section_2: task_1 = DummyOperator(task_id="task_1") with TaskGroup("inner_section_2", tooltip="Tasks for inner_section2"): task_2 = BashOperator(task_id="task_2", bash_command='echo 1') task_3 = DummyOperator(task_id="task_3") task_4 = DummyOperator(task_id="task_4") [task_2, task_3] >> task_4 end = DummyOperator(task_id='end') start >> section_1 >> section_2 >> end dot = dot_renderer.render_dag(dag) assert dot.source == '\n'.join( [ 'digraph example_task_group {', '\tgraph [label=example_task_group labelloc=t rankdir=LR]', '\tend [color="#000000" fillcolor="#e8f7e4" label=end shape=rectangle ' 'style="filled,rounded"]', '\tsubgraph cluster_section_1 {', '\t\tcolor="#000000" fillcolor="#6495ed7f" label=section_1 shape=rectangle style=filled', '\t\t"section_1.upstream_join_id" [color="#000000" fillcolor=CornflowerBlue height=0.2 ' 'label="" shape=circle style="filled,rounded" width=0.2]', '\t\t"section_1.downstream_join_id" [color="#000000" fillcolor=CornflowerBlue height=0.2 ' 'label="" shape=circle style="filled,rounded" width=0.2]', '\t\t"section_1.task_1" [color="#000000" fillcolor="#e8f7e4" label=task_1 shape=rectangle ' 'style="filled,rounded"]', '\t\t"section_1.task_2" [color="#000000" fillcolor="#f0ede4" label=task_2 shape=rectangle ' 'style="filled,rounded"]', '\t\t"section_1.task_3" [color="#000000" fillcolor="#e8f7e4" label=task_3 shape=rectangle ' 'style="filled,rounded"]', '\t}', '\tsubgraph cluster_section_2 {', '\t\tcolor="#000000" fillcolor="#6495ed7f" label=section_2 shape=rectangle style=filled', '\t\t"section_2.upstream_join_id" [color="#000000" fillcolor=CornflowerBlue height=0.2 ' 'label="" shape=circle style="filled,rounded" width=0.2]', '\t\t"section_2.downstream_join_id" [color="#000000" fillcolor=CornflowerBlue height=0.2 ' 'label="" shape=circle style="filled,rounded" width=0.2]', '\t\tsubgraph "cluster_section_2.inner_section_2" {', '\t\t\tcolor="#000000" fillcolor="#6495ed7f" label=inner_section_2 shape=rectangle ' 'style=filled', '\t\t\t"section_2.inner_section_2.task_2" [color="#000000" fillcolor="#f0ede4" label=task_2 ' 'shape=rectangle style="filled,rounded"]', '\t\t\t"section_2.inner_section_2.task_3" [color="#000000" fillcolor="#e8f7e4" label=task_3 ' 'shape=rectangle style="filled,rounded"]', '\t\t\t"section_2.inner_section_2.task_4" [color="#000000" fillcolor="#e8f7e4" label=task_4 ' 'shape=rectangle style="filled,rounded"]', '\t\t}', '\t\t"section_2.task_1" [color="#000000" fillcolor="#e8f7e4" label=task_1 shape=rectangle ' 'style="filled,rounded"]', '\t}', '\tstart [color="#000000" fillcolor="#e8f7e4" label=start shape=rectangle ' 'style="filled,rounded"]', '\t"section_1.downstream_join_id" -> "section_2.upstream_join_id"', '\t"section_1.task_1" -> "section_1.task_2"', '\t"section_1.task_1" -> "section_1.task_3"', '\t"section_1.task_2" -> "section_1.downstream_join_id"', '\t"section_1.task_3" -> "section_1.downstream_join_id"', '\t"section_1.upstream_join_id" -> "section_1.task_1"', '\t"section_2.downstream_join_id" -> end', '\t"section_2.inner_section_2.task_2" -> "section_2.inner_section_2.task_4"', '\t"section_2.inner_section_2.task_3" -> "section_2.inner_section_2.task_4"', '\t"section_2.inner_section_2.task_4" -> "section_2.downstream_join_id"', '\t"section_2.task_1" -> "section_2.downstream_join_id"', '\t"section_2.upstream_join_id" -> "section_2.inner_section_2.task_2"', '\t"section_2.upstream_join_id" -> "section_2.inner_section_2.task_3"', '\t"section_2.upstream_join_id" -> "section_2.task_1"', '\tstart -> "section_1.upstream_join_id"', '}', ] )