From daf8f31080f06c044b4336071bd383bbbcdc6085 Mon Sep 17 00:00:00 2001 From: Tomek Urbaszek Date: Wed, 23 Sep 2020 15:31:40 +0200 Subject: [PATCH] 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 <34898234+olchas@users.noreply.github.com> Co-authored-by: Felix Uellendall --- airflow/models/baseoperator.py | 6 +++- airflow/operators/bash.py | 1 + .../google/cloud/operators/dataproc.py | 1 + airflow/www/utils.py | 10 +++++-- airflow/www/views.py | 9 ++++-- docs/howto/custom-operator.rst | 29 +++++++++++++++++++ tests/serialization/test_dag_serialization.py | 2 ++ 7 files changed, 52 insertions(+), 6 deletions(-) diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py index 7002379a682d..8dcb60f5b5d8 100644 --- 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 index e153a0d4cc88..2d55d3f91ccb 100644 --- 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 index 42f6f147600f..03efe00f1993 100644 --- 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 index 6dc69419afff..4db2dd71f6d7 100644 --- a/airflow/www/utils.py +++ b/airflow/www/utils.py @@ -333,12 +333,12 @@ def wrapped_markdown(s, css_class=None): '
'.format(css_class=css_class) + markdown.markdown(s) + "
" ) -# 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 index cb6d41e5e3fb..a91c4dce04cf 100644 --- 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("
{}
").format(pformat(content)) # noqa diff --git a/docs/howto/custom-operator.rst b/docs/howto/custom-operator.rst index d10800c90675..568799ca63b2 100644 --- 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 index 1b3a9930754c..8d9c9230327e 100644 --- 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",