Add template fields renderers for better UI rendering (apache#11061)

This PR adds possibility to define template_fields_renderers for an operator. In this way users will be able to provide information what lexer should be used for rendering a particular field. This is super useful for custom operator and gives more flexibility than predefined keywords. Co-authored-by: Kamil Olszewski <[email protected]> Co-authored-by: Felix Uellendall <[email protected]>
dropbox · Sep 23, 2020 · daf8f31 · daf8f31
1 parent 0edc3dd
commit daf8f31
Show file tree

Hide file tree

Showing 7 changed files with 52 additions and 6 deletions.
diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py
@@ -276,6 +276,9 @@ class derived from this one results in the creation of a task object,
     template_fields: Iterable[str] = ()
     # Defines which files extensions to look for in the templated fields
     template_ext: Iterable[str] = ()
+    # Template field renderers indicating type of the field, for example sql, json, bash
+    template_fields_renderers: Dict[str, str] = {}
+
     # Defines the color in the UI
     ui_color = '#fff'  # type: str
     ui_fgcolor = '#000'  # type: str
@@ -1311,7 +1314,8 @@ def get_serialized_fields(cls):
                 vars(BaseOperator(task_id='test')).keys() - {
                     'inlets', 'outlets', '_upstream_task_ids', 'default_args', 'dag', '_dag',
                     '_BaseOperator__instantiated',
-                } | {'_task_type', 'subdag', 'ui_color', 'ui_fgcolor', 'template_fields'})
+                } | {'_task_type', 'subdag', 'ui_color', 'ui_fgcolor',
+                     'template_fields', 'template_fields_renderers'})
 
         return cls.__serialized_fields
 

diff --git a/airflow/operators/bash.py b/airflow/operators/bash.py
@@ -95,6 +95,7 @@ class BashOperator(BaseOperator):
     """
 
     template_fields = ('bash_command', 'env')
+    template_fields_renderers = {'bash_command': 'bash', 'env': 'json'}
     template_ext = ('.sh', '.bash',)
     ui_color = '#f0ede4'
 

diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py
@@ -460,6 +460,7 @@ class DataprocCreateClusterOperator(BaseOperator):
         'labels',
         'impersonation_chain',
     )
+    template_fields_renderers = {'cluster_config': 'json'}
 
     @apply_defaults
     def __init__(  # pylint: disable=too-many-arguments

diff --git a/airflow/www/utils.py b/airflow/www/utils.py
@@ -333,12 +333,12 @@ def wrapped_markdown(s, css_class=None):
         '<div class="rich_doc {css_class}" >'.format(css_class=css_class) + markdown.markdown(s) + "</div>"
     )
 
-# pylint: disable=no-member
-
 
+# pylint: disable=no-member
 def get_attr_renderer():
     """Return Dictionary containing different Pygments Lexers for Rendering & Highlighting"""
     return {
+        'bash': lambda x: render(x, lexers.BashLexer),
         'bash_command': lambda x: render(x, lexers.BashLexer),
         'hql': lambda x: render(x, lexers.SqlLexer),
         'sql': lambda x: render(x, lexers.SqlLexer),
@@ -347,9 +347,13 @@ def get_attr_renderer():
         'doc_rst': lambda x: render(x, lexers.RstLexer),
         'doc_yaml': lambda x: render(x, lexers.YamlLexer),
         'doc_md': wrapped_markdown,
+        'json': lambda x: render(x, lexers.JsonLexer),
+        'md': wrapped_markdown,
+        'py': lambda x: render(get_python_source(x), lexers.PythonLexer),
         'python_callable': lambda x: render(get_python_source(x), lexers.PythonLexer),
+        'rst': lambda x: render(x, lexers.RstLexer),
+        'yaml': lambda x: render(x, lexers.YamlLexer),
     }
-
 # pylint: enable=no-member
 
 

diff --git a/airflow/www/views.py b/airflow/www/views.py
@@ -823,10 +823,15 @@ def rendered(self):
             flash("Error rendering template: " + str(e), "error")
         title = "Rendered Template"
         html_dict = {}
+        renderers = wwwutils.get_attr_renderer()
+
         for template_field in task.template_fields:
             content = getattr(task, template_field)
-            if template_field in wwwutils.get_attr_renderer():
-                html_dict[template_field] = wwwutils.get_attr_renderer()[template_field](content)
+            renderer = task.template_fields_renderers.get(template_field, template_field)
+            if renderer in renderers:
+                if isinstance(content, (dict, list)):
+                    content = json.dumps(content, sort_keys=True, indent=4)
+                html_dict[template_field] = renderers[renderer](content)
             else:
                 html_dict[template_field] = \
                     Markup("<pre><code>{}</pre></code>").format(pformat(content))  # noqa

diff --git a/docs/howto/custom-operator.rst b/docs/howto/custom-operator.rst
@@ -200,6 +200,35 @@ with actual value. Note that Jinja substitutes the operator attributes and not t
 
 In the example, the ``template_fields`` should be ``['guest_name']`` and not  ``['name']``
 
+Additionally you may provide ``template_fields_renderers`` dictionary which defines in what style the value
+from template field renders in Web UI. For example:
+
+.. code-block:: python
+
+        class MyRequestOperator(BaseOperator):
+            template_fields = ['request_body']
+            template_fields_renderers = {'request_body': 'json'}
+
+            @apply_defaults
+            def __init__(
+                    self,
+                    request_body: str,
+                    **kwargs) -> None:
+                super().__init__(**kwargs)
+                self.request_body = request_body
+
+Currently available lexers:
+
+  - bash
+  - doc
+  - json
+  - md
+  - py
+  - rst
+  - sql
+  - yaml
+
+If you use a non existing lexer then the value of the template field will be rendered as a pretty printed object.
 
 Define an operator extra link
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py
@@ -93,6 +93,7 @@
                 "ui_color": "#f0ede4",
                 "ui_fgcolor": "#000",
                 "template_fields": ['bash_command', 'env'],
+                "template_fields_renderers": {'bash_command': 'bash', 'env': 'json'},
                 "bash_command": "echo {{ task.task_id }}",
                 'label': 'bash_task',
                 "_task_type": "BashOperator",
@@ -116,6 +117,7 @@
                 "ui_color": "#fff",
                 "ui_fgcolor": "#000",
                 "template_fields": ['bash_command'],
+                "template_fields_renderers": {},
                 "_task_type": "CustomOperator",
                 "_task_module": "tests.test_utils.mock_operators",
                 "pool": "default_pool",