[airflow] branch master updated: Add support for non-default orientation in `dag show` command (#8834)

[ka...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

kamilbregula pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/master by this push:
     new 78a48db  Add support for non-default orientation in `dag show` command (#8834)
78a48db is described below

commit 78a48db75be89d585020124d584a772578f99f57
Author: Kallam Reddy <kl...@gmail.com>
AuthorDate: Mon May 11 20:02:26 2020 -0700

    Add support for non-default orientation in `dag show` command (#8834)
---
 airflow/utils/dot_renderer.py    |  3 ++-
 tests/utils/test_dot_renderer.py | 30 ++++++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/airflow/utils/dot_renderer.py b/airflow/utils/dot_renderer.py
index d5deb9e..605dd82 100644
--- a/airflow/utils/dot_renderer.py
+++ b/airflow/utils/dot_renderer.py
@@ -57,7 +57,8 @@ def render_dag(dag: DAG, tis: Optional[List[TaskInstance]] = None) -> graphviz.D
     :return: Graphviz object
     :rtype: graphviz.Digraph
     """
-    dot = graphviz.Digraph(dag.dag_id, graph_attr={"rankdir": "LR", "labelloc": "t", "label": dag.dag_id})
+    dot = graphviz.Digraph(dag.dag_id, graph_attr={"rankdir": dag.orientation if dag.orientation else "LR",
+                                                   "labelloc": "t", "label": dag.dag_id})
     states_by_task_id = None
     if tis is not None:
         states_by_task_id = {ti.task_id: ti.state for ti in tis}
diff --git a/tests/utils/test_dot_renderer.py b/tests/utils/test_dot_renderer.py
index 041c947..a18c134 100644
--- a/tests/utils/test_dot_renderer.py
+++ b/tests/utils/test_dot_renderer.py
@@ -76,3 +76,33 @@ class TestDotRenderer(unittest.TestCase):
         self.assertIn('first [color=black fillcolor=tan shape=rectangle style="filled,rounded"]', source)
         self.assertIn('second [color=white fillcolor=green shape=rectangle style="filled,rounded"]', source)
         self.assertIn('third [color=black fillcolor=lime shape=rectangle style="filled,rounded"]', source)
+
+    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
+        self.assertIn("label=DAG_ID", source)
+        self.assertIn(f'label=DAG_ID labelloc=t rankdir={orientation}', 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
+        self.assertIn("label=DAG_ID", source)
+        self.assertIn(f'label=DAG_ID labelloc=t rankdir={orientation}', source)