[v2-11-test] Ensuring XCom return value can be mapped for dynamically-mapped @task_group's

airflow/decorators/task_group.py

@@ -38,6 +38,7 @@
     ListOfDictsExpandInput,
     MappedArgument,
 )
+from airflow.models.mappedoperator import ensure_xcomarg_return_value
 from airflow.models.taskmixin import DAGNode
 from airflow.models.xcom_arg import XComArg
 from airflow.typing_compat import ParamSpec
@@ -134,6 +135,11 @@ def expand(self, **kwargs: OperatorExpandArgument) -> DAGNode:
         self._validate_arg_names("expand", kwargs)
         prevent_duplicates(self.partial_kwargs, kwargs, fail_reason="mapping already partial")
         expand_input = DictOfListsExpandInput(kwargs)
+
+        # Similar to @task, @task_group should not be "mappable" over an XCom with a custom key. This will
+        # raise an exception, rather than having an ambiguous exception similar to the one found in #51109.
+        ensure_xcomarg_return_value(expand_input.value)
+
         return self._create_task_group(
             functools.partial(MappedTaskGroup, expand_input=expand_input),
             **self.partial_kwargs,
@@ -163,6 +169,11 @@ def expand_kwargs(self, kwargs: OperatorExpandKwargsArgument) -> DAGNode:
         map_kwargs = (k for k in self.function_signature.parameters if k not in self.partial_kwargs)
 
         expand_input = ListOfDictsExpandInput(kwargs)
+
+        # Similar to @task, @task_group should not be "mappable" over an XCom with a custom key. This will
+        # raise an exception, rather than having an ambiguous exception similar to the one found in #51109.
+        ensure_xcomarg_return_value(expand_input.value)
+
         return self._create_task_group(
             functools.partial(MappedTaskGroup, expand_input=expand_input),
             **self.partial_kwargs,

tests/decorators/test_task_group.py

@@ -178,6 +178,31 @@ def tg(a, b):
     assert saved == {"a": 1, "b": MappedArgument(input=tg._expand_input, key="b")}
 
 
+def test_expand_invalid_xcomarg_return_value():
+    saved = {}
+
+    @dag(schedule=None, start_date=pendulum.datetime(2022, 1, 1))
+    def pipeline():
+        @task
+        def t():
+            return {"values": ["value_1", "value_2"]}
+
+        @task_group()
+        def tg(a, b):
+            saved["a"] = a
+            saved["b"] = b
+
+        tg.partial(a=1).expand(b=t()["values"])
+
+    with pytest.raises(ValueError) as ctx:
+        pipeline()
+
+    assert (
+        str(ctx.value)
+        == "cannot map over XCom with custom key 'values' from <Task(_PythonDecoratedOperator): t>"
+    )
+
+
 def test_expand_kwargs_no_wildcard():
     @dag(schedule=None, start_date=pendulum.datetime(2022, 1, 1))
     def pipeline():
@@ -262,6 +287,31 @@ def t2():
     assert "missing upstream values: ['b']" not in caplog.text
 
 
+def test_expand_kwargs_invalid_xcomarg_return_value():
+    saved = {}
+
+    @dag(schedule=None, start_date=pendulum.datetime(2022, 1, 1))
+    def pipeline():
+        @task
+        def t():
+            return {"values": [{"b": 2}, {"b": 3}]}
+
+        @task_group()
+        def tg(a, b):
+            saved["a"] = a
+            saved["b"] = b
+
+        tg.partial(a=1).expand_kwargs(t()["values"])
+
+    with pytest.raises(ValueError) as ctx:
+        pipeline()
+
+    assert (
+        str(ctx.value)
+        == "cannot map over XCom with custom key 'values' from <Task(_PythonDecoratedOperator): t>"
+    )
+
+
 def test_override_dag_default_args():
     @dag(
         dag_id="test_dag",