diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py index 66980aee35cd2..4a0f6c007adc1 100644 --- a/airflow/providers/google/cloud/operators/bigquery.py +++ b/airflow/providers/google/cloud/operators/bigquery.py @@ -2196,7 +2196,7 @@ class BigQueryInsertJobOperator(BaseOperator): "impersonation_chain", ) template_ext = (".json",) - template_fields_renderers = {"configuration": "json"} + template_fields_renderers = {"configuration": "json", "configuration.query.query": "sql"} ui_color = BigQueryUIColors.QUERY.value def __init__( diff --git a/airflow/www/views.py b/airflow/www/views.py index 3769c1ac4f072..fdde56fc750be 100644 --- a/airflow/www/views.py +++ b/airflow/www/views.py @@ -279,6 +279,29 @@ def task_group_to_dict(task_group): } +def get_key_paths(input_dict): + """Return a list of dot-separated dictionary paths""" + for key, value in input_dict.items(): + if isinstance(value, dict): + for sub_key in get_key_paths(value): + yield '.'.join((key, sub_key)) + else: + yield key + + +def get_value_from_path(key_path, content): + """Return the value from a dictionary based on dot-separated path of keys""" + elem = content + for x in key_path.strip(".").split("."): + try: + x = int(x) + elem = elem[x] + except ValueError: + elem = elem.get(x) + + return elem + + def dag_edges(dag): """ Create the list of edges needed to construct the Graph view. @@ -995,11 +1018,31 @@ def rendered_templates(self): 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) + json_content = json.dumps(content, sort_keys=True, indent=4) + html_dict[template_field] = renderers[renderer](json_content) + else: + html_dict[template_field] = renderers[renderer](content) else: html_dict[template_field] = Markup("
{}
").format(pformat(content)) + if isinstance(content, dict): + if template_field == 'op_kwargs': + for key, value in content.items(): + renderer = task.template_fields_renderers.get(key, key) + if renderer in renderers: + html_dict['.'.join([template_field, key])] = renderers[renderer](value) + else: + html_dict['.'.join([template_field, key])] = Markup( + "
{}
" + ).format(pformat(value)) + else: + for dict_keys in get_key_paths(content): + template_path = '.'.join((template_field, dict_keys)) + renderer = task.template_fields_renderers.get(template_path, template_path) + if renderer in renderers: + content_value = get_value_from_path(dict_keys, content) + html_dict[template_path] = renderers[renderer](content_value) + return self.render_template( 'airflow/ti_code.html', html_dict=html_dict, diff --git a/docs/apache-airflow/howto/custom-operator.rst b/docs/apache-airflow/howto/custom-operator.rst index 8074bf534eb0a..25558f76458f5 100644 --- a/docs/apache-airflow/howto/custom-operator.rst +++ b/docs/apache-airflow/howto/custom-operator.rst @@ -195,7 +195,7 @@ 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 +Additionally you may provide ``template_fields_renderers`` a dictionary which defines in what style the value from template field renders in Web UI. For example: .. code-block:: python @@ -208,12 +208,48 @@ from template field renders in Web UI. For example: super().__init__(**kwargs) self.request_body = request_body +In the situation where ``template_field`` is itself a dictionary, it is also possible to specify a +dot-separated key path to extract and render individual elements appropriately. For example: + +.. code-block:: python + + class MyConfigOperator(BaseOperator): + template_fields = ["configuration"] + template_fields_renderers = { + "configuration": "json", + "configuration.query.sql": "sql", + } + + def __init__(self, configuration: dict, **kwargs) -> None: + super().__init__(**kwargs) + self.configuration = configuration + +Then using this template as follows: + +.. code-block:: python + + with dag: + config_task = MyConfigOperator( + task_id="task_id_1", + configuration={"query": {"job_id": "123", "sql": "select * from my_table"}}, + dag=dag, + ) + +This will result in the UI rendering ``configuration`` as json in addition to the value contained in the +configuration at ``query.sql`` to be rendered with the SQL lexer. + +.. image:: ../img/template_field_renderer_path.png + Currently available lexers: - bash - doc + - hql + - html + - jinja - json - md + - powershell - py - rst - sql diff --git a/docs/apache-airflow/img/template_field_renderer_path.png b/docs/apache-airflow/img/template_field_renderer_path.png new file mode 100644 index 0000000000000..bbecf61ea5da0 Binary files /dev/null and b/docs/apache-airflow/img/template_field_renderer_path.png differ diff --git a/tests/www/views/test_views.py b/tests/www/views/test_views.py index ebb4f462a2985..a15f708e5f354 100644 --- a/tests/www/views/test_views.py +++ b/tests/www/views/test_views.py @@ -22,7 +22,7 @@ from airflow.configuration import initialize_config from airflow.plugins_manager import AirflowPlugin, EntryPointSource -from airflow.www.views import get_safe_url, truncate_task_duration +from airflow.www.views import get_key_paths, get_safe_url, get_value_from_path, truncate_task_duration from tests.test_utils.config import conf_vars from tests.test_utils.mock_plugins import mock_plugin_manager from tests.test_utils.www import check_content_in_response, check_content_not_in_response @@ -243,3 +243,26 @@ def get_task_instance(session, task): dagrun.refresh_from_db(session=session) # dagrun should be set to QUEUED assert dagrun.get_state() == State.QUEUED + + +TEST_CONTENT_DICT = {"key1": {"key2": "val2", "key3": "val3", "key4": {"key5": "val5"}}} + + +@pytest.mark.parametrize( + "test_content_dict, expected_paths", [(TEST_CONTENT_DICT, ("key1.key2", "key1.key3", "key1.key4.key5"))] +) +def test_generate_key_paths(test_content_dict, expected_paths): + for key_path in get_key_paths(test_content_dict): + assert key_path in expected_paths + + +@pytest.mark.parametrize( + "test_content_dict, test_key_path, expected_value", + [ + (TEST_CONTENT_DICT, "key1.key2", "val2"), + (TEST_CONTENT_DICT, "key1.key3", "val3"), + (TEST_CONTENT_DICT, "key1.key4.key5", "val5"), + ], +) +def test_get_value_from_path(test_content_dict, test_key_path, expected_value): + assert expected_value == get_value_from_path(test_key_path, test_content_dict)