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)