[airflow] branch main updated: Only excluded actually expanded fields from render (#25599)

[ephraimanierobi]

This is an automated email from the ASF dual-hosted git repository.

ephraimanierobi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new ed39703cd4 Only excluded actually expanded fields from render (#25599)
ed39703cd4 is described below

commit ed39703cd4f619104430b91d7ba67f261e5bfddb
Author: Tzu-ping Chung <uranu...@gmail.com>
AuthorDate: Mon Aug 15 20:02:30 2022 +0800

    Only excluded actually expanded fields from render (#25599)
    
    Previously we excluded all mapped kwargs from being rendered, which
    cause issues for fields that are not actually rendered during expansion
    (i.e. those that are literals). This new implementation adds some extra
    logic to ensure we still include kwargs that still need rendering.
---
 airflow/models/mappedoperator.py    |  9 +++++----
 tests/models/test_mappedoperator.py | 18 ++++++++++++++++++
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py
index f62d1c687b..2a93e442b7 100644
--- a/airflow/models/mappedoperator.py
+++ b/airflow/models/mappedoperator.py
@@ -730,10 +730,11 @@ class MappedOperator(AbstractOperator):
             return None
 
     def _get_template_fields_to_render(self, expanded: Iterable[str]) -> 
Iterable[str]:
-        # Since the mapped kwargs are already resolved during unmapping,
-        # they must be removed from the list of templated fields to avoid
-        # being rendered again (which breaks escaping).
-        return set(self.template_fields).difference(expanded)
+        # Mapped kwargs from XCom are already resolved during unmapping, so 
they
+        # must be removed from the list of templated fields to avoid being
+        # rendered again.
+        unexpanded_keys = {k for k, _ in 
self._get_specified_expand_input().iter_parse_time_resolved_kwargs()}
+        return set(self.template_fields).difference(k for k in expanded if k 
not in unexpanded_keys)
 
     def render_template_fields(
         self,
diff --git a/tests/models/test_mappedoperator.py 
b/tests/models/test_mappedoperator.py
index 483ec97c54..cf80e6ede4 100644
--- a/tests/models/test_mappedoperator.py
+++ b/tests/models/test_mappedoperator.py
@@ -313,6 +313,24 @@ def 
test_mapped_render_template_fields_validating_operator(dag_maker, session):
     assert op.arg2 == "a"
 
 
+def test_mapped_render_nested_template_fields(dag_maker, session):
+    with dag_maker(session=session):
+        MockOperator.partial(task_id="t").expand(arg1=["{{ ti.task_id }}", 
["s", "{{ ti.task_id }}"]])
+
+    dr = dag_maker.create_dagrun()
+    decision = dr.task_instance_scheduling_decisions()
+    tis = {(ti.task_id, ti.map_index): ti for ti in decision.schedulable_tis}
+    assert len(tis) == 2
+
+    ti = tis[("t", 0)]
+    ti.run(session=session)
+    assert ti.task.arg1 == "t"
+
+    ti = tis[("t", 1)]
+    ti.run(session=session)
+    assert ti.task.arg1 == ["s", "t"]
+
+
 @pytest.mark.parametrize(
     ["num_existing_tis", "expected"],
     (