Add template fields renderers for better UI rendering (#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]>
apache · 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",