Allow changing Task States Colors (#9520)

(cherry picked from commit e1108d42d2b3b9c049ed8872176923621b4c8d79) GitOrigin-RevId: 80089bab2ea1d2a29b678e813ea026d841c4cfde
GoogleCloudPlatform · Jun 4, 2021 · dd10cee · dd10cee
1 parent 2ac1b12
commit dd10cee
Show file tree

Hide file tree

Showing 14 changed files with 97 additions and 54 deletions.
diff --git a/airflow/settings.py b/airflow/settings.py
@@ -153,6 +153,20 @@ def timing(self, stat, dt):
 # The JSON library to use for DAG Serialization and De-Serialization
 json = json
 
+# Dictionary containing State and colors associated to each state to
+# display on the Webserver
+STATE_COLORS = {
+    "queued": "gray",
+    "running": "lime",
+    "success": "green",
+    "failed": "red",
+    "up_for_retry": "gold",
+    "up_for_reschedule": "turquoise",
+    "upstream_failed": "orange",
+    "skipped": "pink",
+    "scheduled": "tan",
+}
+
 
 def policy(task_instance):
     """

diff --git a/airflow/utils/state.py b/airflow/utils/state.py
@@ -21,6 +21,8 @@
 
 from builtins import object
 
+from airflow.settings import STATE_COLORS
+
 
 class State(object):
     """
@@ -80,6 +82,7 @@ class State(object):
         SCHEDULED: 'tan',
         NONE: 'lightblue',
     }
+    state_color.update(STATE_COLORS)  # type: ignore
 
     @classmethod
     def color(cls, state):

diff --git a/airflow/www/app.py b/airflow/www/app.py
@@ -32,7 +32,7 @@
 from airflow import models, version, LoggingMixin
 from airflow.configuration import conf
 from airflow.models.connection import Connection
-from airflow.settings import Session
+from airflow.settings import Session, STATE_COLORS
 
 from airflow.www.blueprints import routes
 from airflow.logging_config import configure_logging
@@ -188,7 +188,8 @@ def jinja_globals():
                 'log_auto_tailing_offset': conf.getint(
                     'webserver', 'log_auto_tailing_offset', fallback=30),
                 'log_animation_speed': conf.getint(
-                    'webserver', 'log_animation_speed', fallback=1000)
+                    'webserver', 'log_animation_speed', fallback=1000),
+                'state_color_mapping': STATE_COLORS
             }
 
         @app.before_request

diff --git a/airflow/www/templates/admin/master.html b/airflow/www/templates/admin/master.html
@@ -23,6 +23,13 @@
   <link href="{{ url_for('static', filename='bootstrap-theme.css') }}" rel="stylesheet">
   <link rel="icon" type="image/png" href="{{ url_for("static", filename="pin_32.png") }}">
   <link rel="stylesheet" type="text/css" href="{{ url_for("static", filename="main.css") }}">
+    <style type="text/css">
+    {% for state, state_color in state_color_mapping.items() %}
+      span.{{state}} {
+        background-color: {{state_color}};
+      }
+    {% endfor %}
+  </style>
 {% endblock %}
 
 {% block tail_js %}

diff --git a/airflow/www/templates/airflow/gantt.html b/airflow/www/templates/airflow/gantt.html
@@ -24,6 +24,13 @@
 <link href="{{ admin_static.url(filename='vendor/bootstrap-daterangepicker/daterangepicker-bs2.css') }}" rel="stylesheet"/>
 <link type="text/css" href="{{ url_for('static', filename='gantt.css') }}" rel="stylesheet" />
 <link type="text/css" href="{{ url_for('static', filename='tree.css') }}" rel="stylesheet" />
+<style type="text/css">
+  {% for state, state_color in state_color_mapping.items() %}
+    rect.{{state}} {
+      fill: {{state_color}};
+    }
+  {% endfor %}
+</style>
 {% endblock %}
 
 {% block body %}

diff --git a/airflow/www/templates/airflow/graph.html b/airflow/www/templates/airflow/graph.html
@@ -27,6 +27,13 @@
 {{ super() }}
 <link rel="stylesheet" type="text/css" href="{{ url_for('static', filename='dagre.css') }}">
 <link rel="stylesheet" type="text/css" href="{{ url_for('static', filename='graph.css') }}">
+<style type="text/css">
+  {% for state, state_color in state_color_mapping.items() %}
+    g.node.{{state}} rect {
+      stroke: {{state_color}};
+    }
+  {% endfor %}
+</style>
 {% endblock %}
 
 {% block body %}
@@ -64,14 +71,9 @@
 
   <div>
     <div class="legend_item state" style="border-color:white;">no_status</div>
-    <div class="legend_item state" style="border-color:grey;">queued</div>
-    <div class="legend_item state" style="border-color:gold;">up_for_retry</div>
-    <div class="legend_item state" style="border-color:turquoise;">up_for_reschedule</div>
-    <div class="legend_item state" style="border-color:orange;">upstream_failed</div>
-    <div class="legend_item state" style="border-color:pink;">skipped</div>
-    <div class="legend_item state" style="border-color:red;">failed</div>
-    <div class="legend_item state" style="border-color:lime;">running</div>
-    <div class="legend_item state" style="border-color:green;">success</div>
+    {% for state, state_color in state_color_mapping.items() %}
+      <div class="legend_item state" style="border-color:{{state_color}};">{{state}}</div>
+    {% endfor %}
   </div>
   <div style="clear:both;"></div>
 </div>

diff --git a/airflow/www/templates/airflow/tree.html b/airflow/www/templates/airflow/tree.html
@@ -24,6 +24,13 @@
 {{ super() }}
 <link rel="stylesheet" type="text/css" href="{{ url_for('static', filename='tree.css') }}">
 <link href="{{ admin_static.url(filename='vendor/bootstrap-daterangepicker/daterangepicker-bs2.css') }}" rel="stylesheet">
+<style type="text/css">
+  {% for state, state_color in state_color_mapping.items() %}
+    rect.{{state}} {
+      fill: {{state_color}};
+    }
+  {% endfor %}
+</style>
 {% endblock %}
 
 {% block body %}
@@ -46,22 +53,10 @@
 <div>
   <div class="legend_item" style="border: none;">no_status</div>
   <div class="square" style="background: white;"></div>
-  <div class="legend_item" style="border: none;">queued</div>
-  <div class="square" style="background: grey;"></div>
-  <div class="legend_item" style="border: none;">up_for_retry</div>
-  <div class="square" style="background: gold;"></div>
-  <div class="legend_item" style="border: none;">up_for_reschedule</div>
-  <div class="square" style="background: turquoise;"></div>
-  <div class="legend_item" style="border: none;">upstream_failed</div>
-  <div class="square" style="background: orange;"></div>
-  <div class="legend_item" style="border: none;">skipped</div>
-  <div class="square" style="background: pink;"></div>
-  <div class="legend_item" style="border: none;">failed</div>
-  <div class="square" style="background: red;"></div>
-  <div class="legend_item" style="border: none;">running</div>
-  <div class="square" style="background: lime;"></div>
-  <div class="legend_item" style="border: none;">success</div>
-  <div class="square" style="background: green;"></div>
+  {% for state, state_color in state_color_mapping.items() %}
+    <div class="legend_item" style="border: none;">{{state}}</div>
+    <div class="square" style="background: {{state_color}};"></div>
+  {% endfor %}
   {% for op in operators %}
   <div class="legend_circle" style="background:{{ op.ui_color }};">
   </div>

diff --git a/airflow/www/views.py b/airflow/www/views.py
@@ -75,7 +75,7 @@
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator, Connection, DagRun, errors, XCom
 from airflow.models.dagcode import DagCode
-from airflow.settings import STORE_SERIALIZED_DAGS
+from airflow.settings import STATE_COLORS, STORE_SERIALIZED_DAGS
 from airflow.operators.subdag_operator import SubDagOperator
 from airflow.ti_deps.dep_context import RUNNING_DEPS, SCHEDULER_QUEUED_DEPS, DepContext
 from airflow.utils import timezone
@@ -2319,6 +2319,7 @@ def get_int_arg(value, default=0):
 
         state_color_mapping = State.state_color.copy()
         state_color_mapping["null"] = state_color_mapping.pop(None)
+        state_color_mapping.update(STATE_COLORS)
 
         return self.render(
             'airflow/dags.html',

diff --git a/airflow/www_rbac/app.py b/airflow/www_rbac/app.py
@@ -35,6 +35,7 @@
 from airflow import settings, version
 from airflow.configuration import conf
 from airflow.logging_config import configure_logging
+from airflow.settings import STATE_COLORS
 from airflow.www_rbac.static_config import configure_manifest_files
 
 app = None  # type: Any
@@ -242,7 +243,8 @@ def jinja_globals():  # pylint: disable=unused-variable
                 'log_auto_tailing_offset': conf.getint(
                     'webserver', 'log_auto_tailing_offset', fallback=30),
                 'log_animation_speed': conf.getint(
-                    'webserver', 'log_animation_speed', fallback=1000)
+                    'webserver', 'log_animation_speed', fallback=1000),
+                'state_color_mapping': STATE_COLORS
             }
 
             if 'analytics_tool' in conf.getsection('webserver'):

diff --git a/airflow/www_rbac/templates/airflow/gantt.html b/airflow/www_rbac/templates/airflow/gantt.html
@@ -21,6 +21,13 @@
 {{ super() }}
 <link type="text/css" href="{{ url_for('static', filename='css/gantt.css') }}" rel="stylesheet" />
 <link type="text/css" href="{{ url_for('static', filename='css/tree.css') }}" rel="stylesheet" />
+<style type="text/css">
+  {% for state, state_color in state_color_mapping.items() %}
+    rect.{{state}} {
+      fill: {{state_color}};
+    }
+  {% endfor %}
+</style>
 {% endblock %}
 
 {% block content %}

diff --git a/airflow/www_rbac/templates/airflow/graph.html b/airflow/www_rbac/templates/airflow/graph.html
@@ -22,6 +22,13 @@
 {% block head_css %}
 {{ super() }}
 <link rel="stylesheet" type="text/css" href="{{ url_for('static', filename='css/graph.css') }}">
+<style type="text/css">
+  {% for state, state_color in state_color_mapping.items() %}
+    g.node.{{state}} rect {
+      stroke: {{state_color}};
+    }
+  {% endfor %}
+</style>
 {% endblock %}
 
 {% block content %}
@@ -58,14 +65,9 @@
 
   <div>
     <div class="legend_item state" style="border-color:white;">no_status</div>
-    <div class="legend_item state" style="border-color:grey;">queued</div>
-    <div class="legend_item state" style="border-color:gold;">up_for_retry</div>
-    <div class="legend_item state" style="border-color:turquoise;">up_for_reschedule</div>
-    <div class="legend_item state" style="border-color:orange;">upstream_failed</div>
-    <div class="legend_item state" style="border-color:pink;">skipped</div>
-    <div class="legend_item state" style="border-color:red;">failed</div>
-    <div class="legend_item state" style="border-color:lime;">running</div>
-    <div class="legend_item state" style="border-color:green;">success</div>
+  {% for state, state_color in state_color_mapping.items() %}
+    <div class="legend_item state" style="border-color:{{state_color}};">{{state}}</div>
+  {% endfor %}
   </div>
   <div style="clear:both;"></div>
 </div>

diff --git a/airflow/www_rbac/templates/airflow/master.html b/airflow/www_rbac/templates/airflow/master.html
@@ -27,7 +27,13 @@
   {% endif %}
   <link href="{{ url_for_asset('main.css') }}" rel="stylesheet">
   <link href="{{ url_for_asset('bootstrap-datetimepicker.min.css') }}" rel="stylesheet">
-
+  <style type="text/css">
+  {% for state, state_color in state_color_mapping.items() %}
+    span.{{state}} {
+      background-color: {{state_color}};
+    }
+  {% endfor %}
+  </style>
   <link rel="icon" type="image/png" href="{{ url_for('static', filename='pin_32.png') }}">
 {% endblock %}
 

diff --git a/airflow/www_rbac/templates/airflow/tree.html b/airflow/www_rbac/templates/airflow/tree.html
@@ -21,6 +21,13 @@
 {% block head_css %}
 {{ super() }}
 <link rel="stylesheet" type="text/css" href="{{ url_for('static', filename='css/tree.css') }}">
+<style type="text/css">
+  {% for state, state_color in state_color_mapping.items() %}
+    rect.{{state}} {
+      fill: {{state_color}};
+    }
+  {% endfor %}
+</style>
 {% endblock %}
 
 {% block content %}
@@ -44,22 +51,10 @@
 <div>
   <div class="legend_item" style="border: none;">no_status</div>
   <div class="square" style="background: white;"></div>
-  <div class="legend_item" style="border: none;">queued</div>
-  <div class="square" style="background: grey;"></div>
-  <div class="legend_item" style="border: none;">up_for_retry</div>
-  <div class="square" style="background: gold;"></div>
-  <div class="legend_item" style="border: none;">up_for_reschedule</div>
-  <div class="square" style="background: turquoise;"></div>
-  <div class="legend_item" style="border: none;">upstream_failed</div>
-  <div class="square" style="background: orange;"></div>
-  <div class="legend_item" style="border: none;">skipped</div>
-  <div class="square" style="background: pink;"></div>
-  <div class="legend_item" style="border: none;">failed</div>
-  <div class="square" style="background: red;"></div>
-  <div class="legend_item" style="border: none;">running</div>
-  <div class="square" style="background: lime;"></div>
-  <div class="legend_item" style="border: none;">success</div>
-  <div class="square" style="background: green;"></div>
+  {% for state, state_color in state_color_mapping.items() %}
+    <div class="legend_item" style="border: none;">{{state}}</div>
+    <div class="square" style="background: {{state_color}};"></div>
+  {% endfor %}
   {% for op in operators %}
   <div class="legend_circle" style="background:{{ op.ui_color }};">
   </div>

diff --git a/airflow/www_rbac/views.py b/airflow/www_rbac/views.py
@@ -63,7 +63,7 @@
 from airflow.models import Connection, DagModel, DagRun, DagTag, Log, SlaMiss, TaskFail, XCom, errors
 from airflow.exceptions import AirflowException
 from airflow.models.dagcode import DagCode
-from airflow.settings import STORE_SERIALIZED_DAGS
+from airflow.settings import STATE_COLORS, STORE_SERIALIZED_DAGS
 from airflow.ti_deps.dep_context import RUNNING_DEPS, SCHEDULER_QUEUED_DEPS, DepContext
 from airflow.utils import timezone
 from airflow.utils.dates import infer_time_unit, scale_time_units
@@ -334,6 +334,7 @@ def get_int_arg(value, default=0):
 
         state_color_mapping = State.state_color.copy()
         state_color_mapping["null"] = state_color_mapping.pop(None)
+        state_color_mapping.update(STATE_COLORS)
 
         return self.render_template(
             'airflow/dags.html',