From 2a55ffe0cd2c04fde2eecdbeff2fa727067069e1 Mon Sep 17 00:00:00 2001
From: Tao feng <tfeng@lyft.com>
Date: Wed, 9 May 2018 08:45:06 -0700
Subject: [PATCH] [AIRFLOW-2086][AIRFLOW-2393] Customize default dagrun number
 in tree view

Closes #3279 from feng-tao/reduce-tree-view

This introduces a new configuration variable to set the default
number of dag runs displayed in the tree view. For large DAGs, this
could cause timeouts in the webserver.
---
 UPDATING.md                                  |  3 +++
 airflow/config_templates/default_airflow.cfg |  3 +++
 airflow/www/templates/airflow/dag.html       |  2 +-
 airflow/www/templates/airflow/dags.html      |  2 +-
 airflow/www/templates/airflow/list_dags.html |  2 +-
 airflow/www/views.py                         | 19 ++++++++++++-------
 airflow/www_rbac/templates/airflow/dag.html  |  2 +-
 airflow/www_rbac/templates/airflow/dags.html |  2 +-
 airflow/www_rbac/views.py                    | 14 +++++++++-----
 9 files changed, 32 insertions(+), 17 deletions(-)

diff --git a/UPDATING.md b/UPDATING.md
index defd95be0c4dc..c9e1395c52c1c 100644
--- a/UPDATING.md
+++ b/UPDATING.md
@@ -5,6 +5,9 @@ assists users migrating to a new version.
 
 ## Airflow Master
 
+### Add a configuration variable(default_dag_run_display_number) to control numbers of dag run for display
+Add a configuration variable(default_dag_run_display_number) under webserver section to control num of dag run to show in UI.
+
 ### Default executor for SubDagOperator is changed to SequentialExecutor
 
 ### New Webserver UI with Role-Based Access Control
diff --git a/airflow/config_templates/default_airflow.cfg b/airflow/config_templates/default_airflow.cfg
index b91961eda6761..33b99ff3520d2 100644
--- a/airflow/config_templates/default_airflow.cfg
+++ b/airflow/config_templates/default_airflow.cfg
@@ -280,6 +280,9 @@ rbac = False
 # Define the color of navigation bar
 navbar_color = #007A87
 
+# Default dagrun to show in UI
+default_dag_run_display_number = 25
+
 
 [email]
 email_backend = airflow.utils.email.send_email_smtp
diff --git a/airflow/www/templates/airflow/dag.html b/airflow/www/templates/airflow/dag.html
index ed84f2715597d..18242d30c640d 100644
--- a/airflow/www/templates/airflow/dag.html
+++ b/airflow/www/templates/airflow/dag.html
@@ -57,7 +57,7 @@ <h4 class="pull-right">
       <li><a href="{{ url_for("airflow.graph", dag_id=dag.dag_id, root=root, execution_date=execution_date) }}">
           <span class="glyphicon glyphicon-certificate" aria-hidden="true"></span>
         Graph View</a></li>
-      <li><a href="{{ url_for("airflow.tree", dag_id=dag.dag_id, num_runs=25, root=root) }}">
+      <li><a href="{{ url_for("airflow.tree", dag_id=dag.dag_id, num_runs=num_runs, root=root) }}">
           <span class="glyphicon glyphicon-tree-deciduous" aria-hidden="true"></span>
         Tree View
       </a></li>
diff --git a/airflow/www/templates/airflow/dags.html b/airflow/www/templates/airflow/dags.html
index d22bfb36eee19..23978905780ea 100644
--- a/airflow/www/templates/airflow/dags.html
+++ b/airflow/www/templates/airflow/dags.html
@@ -145,7 +145,7 @@ <h2>DAGs</h2>
                 </a>
 
                 <!-- Tree -->
-                <a href="{{ url_for('airflow.tree', dag_id=dag.dag_id, num_runs=25) }}">
+                <a href="{{ url_for('airflow.tree', dag_id=dag.dag_id, num_runs=num_runs) }}">
                     <span class="glyphicon glyphicon-tree-deciduous" aria-hidden="true" data-original-title="Tree View"></span>
                 </a>
 
diff --git a/airflow/www/templates/airflow/list_dags.html b/airflow/www/templates/airflow/list_dags.html
index e8533d7f8297d..c7f2497b94735 100644
--- a/airflow/www/templates/airflow/list_dags.html
+++ b/airflow/www/templates/airflow/list_dags.html
@@ -147,7 +147,7 @@ <h2>DAGs</h2>
                 <svg height="10" width="10" id='dag-{{ row.dag_id }}' style="display: block;"></svg>
             </td>
             <td>
-              <a href="{{ url_for("airflow.tree", dag_id=row.dag_id, num_runs=25) }}" title="Tree View">
+              <a href="{{ url_for("airflow.tree", dag_id=row.dag_id, num_runs=num_runs) }}" title="Tree View">
                   <span class="glyphicon glyphicon-tree-deciduous" aria-hidden="true"></span>
               </a>
               <a href="{{ url_for("airflow.graph", dag_id=row.dag_id) }}" title="Graph View">
diff --git a/airflow/www/views.py b/airflow/www/views.py
index 6e2f1fc6a8f6c..c36d55f270798 100644
--- a/airflow/www/views.py
+++ b/airflow/www/views.py
@@ -7,9 +7,9 @@
 # to you under the Apache License, Version 2.0 (the
 # "License"); you may not use this file except in compliance
 # with the License.  You may obtain a copy of the License at
-# 
+#
 #   http://www.apache.org/licenses/LICENSE-2.0
-# 
+#
 # Unless required by applicable law or agreed to in writing,
 # software distributed under the License is distributed on an
 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
@@ -1227,6 +1227,7 @@ def success(self):
     @wwwutils.action_logging
     @provide_session
     def tree(self, session=None):
+        default_dag_run = conf.getint('webserver', 'default_dag_run_display_number')
         dag_id = request.args.get('dag_id')
         blur = conf.getboolean('webserver', 'demo_mode')
         dag = dagbag.get_dag(dag_id)
@@ -1239,7 +1240,7 @@ def tree(self, session=None):
 
         base_date = request.args.get('base_date')
         num_runs = request.args.get('num_runs')
-        num_runs = int(num_runs) if num_runs else 25
+        num_runs = int(num_runs) if num_runs else default_dag_run
 
         if base_date:
             base_date = timezone.parse(base_date)
@@ -1348,7 +1349,7 @@ def set_duration(tid):
             ),
             root=root,
             form=form,
-            dag=dag, data=data, blur=blur)
+            dag=dag, data=data, blur=blur, num_runs=num_runs)
 
     @expose('/graph')
     @login_required
@@ -1468,11 +1469,12 @@ class GraphForm(Form):
     @wwwutils.action_logging
     @provide_session
     def duration(self, session=None):
+        default_dag_run = conf.getint('webserver', 'default_dag_run_display_number')
         dag_id = request.args.get('dag_id')
         dag = dagbag.get_dag(dag_id)
         base_date = request.args.get('base_date')
         num_runs = request.args.get('num_runs')
-        num_runs = int(num_runs) if num_runs else 25
+        num_runs = int(num_runs) if num_runs else default_dag_run
 
         if base_date:
             base_date = pendulum.parse(base_date)
@@ -1575,11 +1577,12 @@ def duration(self, session=None):
     @wwwutils.action_logging
     @provide_session
     def tries(self, session=None):
+        default_dag_run = conf.getint('webserver', 'default_dag_run_display_number')
         dag_id = request.args.get('dag_id')
         dag = dagbag.get_dag(dag_id)
         base_date = request.args.get('base_date')
         num_runs = request.args.get('num_runs')
-        num_runs = int(num_runs) if num_runs else 25
+        num_runs = int(num_runs) if num_runs else default_dag_run
 
         if base_date:
             base_date = pendulum.parse(base_date)
@@ -1638,11 +1641,12 @@ def tries(self, session=None):
     @wwwutils.action_logging
     @provide_session
     def landing_times(self, session=None):
+        default_dag_run = conf.getint('webserver', 'default_dag_run_display_number')
         dag_id = request.args.get('dag_id')
         dag = dagbag.get_dag(dag_id)
         base_date = request.args.get('base_date')
         num_runs = request.args.get('num_runs')
-        num_runs = int(num_runs) if num_runs else 25
+        num_runs = int(num_runs) if num_runs else default_dag_run
 
         if base_date:
             base_date = pendulum.parse(base_date)
@@ -1759,6 +1763,7 @@ def refresh_all(self):
     @wwwutils.action_logging
     @provide_session
     def gantt(self, session=None):
+
         dag_id = request.args.get('dag_id')
         dag = dagbag.get_dag(dag_id)
         demo_mode = conf.getboolean('webserver', 'demo_mode')
diff --git a/airflow/www_rbac/templates/airflow/dag.html b/airflow/www_rbac/templates/airflow/dag.html
index 02c2fddd7b7ed..35308bea62a18 100644
--- a/airflow/www_rbac/templates/airflow/dag.html
+++ b/airflow/www_rbac/templates/airflow/dag.html
@@ -56,7 +56,7 @@ <h4 class="pull-right">
       <li><a href="{{ url_for("Airflow.graph", dag_id=dag.dag_id, root=root) }}">
           <span class="glyphicon glyphicon-certificate" aria-hidden="true"></span>
         Graph View</a></li>
-      <li><a href="{{ url_for("Airflow.tree", dag_id=dag.dag_id, num_runs=25, root=root) }}">
+      <li><a href="{{ url_for("Airflow.tree", dag_id=dag.dag_id, num_runs=num_runs, root=root) }}">
           <span class="glyphicon glyphicon-tree-deciduous" aria-hidden="true"></span>
         Tree View
       </a></li>
diff --git a/airflow/www_rbac/templates/airflow/dags.html b/airflow/www_rbac/templates/airflow/dags.html
index fc8418dc742f7..a712e5a9ee4bc 100644
--- a/airflow/www_rbac/templates/airflow/dags.html
+++ b/airflow/www_rbac/templates/airflow/dags.html
@@ -146,7 +146,7 @@ <h2>DAGs</h2>
                 </a>
 
                 <!-- Tree -->
-                <a href="{{ url_for('Airflow.tree', dag_id=dag.dag_id, num_runs=25) }}">
+                <a href="{{ url_for('Airflow.tree', dag_id=dag.dag_id, num_runs=num_runs) }}">
                     <span class="glyphicon glyphicon-tree-deciduous" aria-hidden="true" data-original-title="Tree View"></span>
                 </a>
 
diff --git a/airflow/www_rbac/views.py b/airflow/www_rbac/views.py
index 7c20a6513f085..e77446263f357 100644
--- a/airflow/www_rbac/views.py
+++ b/airflow/www_rbac/views.py
@@ -860,6 +860,7 @@ def success(self):
     @action_logging
     @provide_session
     def tree(self, session=None):
+        default_dag_run = conf.getint('webserver', 'default_dag_run_display_number')
         dag_id = request.args.get('dag_id')
         blur = conf.getboolean('webserver', 'demo_mode')
         dag = dagbag.get_dag(dag_id)
@@ -872,7 +873,7 @@ def tree(self, session=None):
 
         base_date = request.args.get('base_date')
         num_runs = request.args.get('num_runs')
-        num_runs = int(num_runs) if num_runs else 25
+        num_runs = int(num_runs) if num_runs else default_dag_run
 
         if base_date:
             base_date = timezone.parse(base_date)
@@ -980,7 +981,7 @@ def set_duration(tid):
             ),
             root=root,
             form=form,
-            dag=dag, data=data, blur=blur)
+            dag=dag, data=data, blur=blur, num_runs=num_runs)
 
     @expose('/graph')
     @has_access
@@ -1101,11 +1102,12 @@ class GraphForm(Form):
     @action_logging
     @provide_session
     def duration(self, session=None):
+        default_dag_run = conf.getint('webserver', 'default_dag_run_display_number')
         dag_id = request.args.get('dag_id')
         dag = dagbag.get_dag(dag_id)
         base_date = request.args.get('base_date')
         num_runs = request.args.get('num_runs')
-        num_runs = int(num_runs) if num_runs else 25
+        num_runs = int(num_runs) if num_runs else default_dag_run
 
         if base_date:
             base_date = pendulum.parse(base_date)
@@ -1204,11 +1206,12 @@ def duration(self, session=None):
     @action_logging
     @provide_session
     def tries(self, session=None):
+        default_dag_run = conf.getint('webserver', 'default_dag_run_display_number')
         dag_id = request.args.get('dag_id')
         dag = dagbag.get_dag(dag_id)
         base_date = request.args.get('base_date')
         num_runs = request.args.get('num_runs')
-        num_runs = int(num_runs) if num_runs else 25
+        num_runs = int(num_runs) if num_runs else default_dag_run
 
         if base_date:
             base_date = pendulum.parse(base_date)
@@ -1267,11 +1270,12 @@ def tries(self, session=None):
     @action_logging
     @provide_session
     def landing_times(self, session=None):
+        default_dag_run = conf.getint('webserver', 'default_dag_run_display_number')
         dag_id = request.args.get('dag_id')
         dag = dagbag.get_dag(dag_id)
         base_date = request.args.get('base_date')
         num_runs = request.args.get('num_runs')
-        num_runs = int(num_runs) if num_runs else 25
+        num_runs = int(num_runs) if num_runs else default_dag_run
 
         if base_date:
             base_date = pendulum.parse(base_date)