Fix mapped tasks does not validate task_id

task-sdk/src/airflow/sdk/bases/decorator.py

@@ -45,6 +45,7 @@
     ListOfDictsExpandInput,
     is_mappable,
 )
+from airflow.sdk.definitions._internal.node import validate_key
 from airflow.sdk.definitions._internal.types import NOTSET
 from airflow.sdk.definitions.asset import Asset
 from airflow.sdk.definitions.context import KNOWN_CONTEXT_KEYS
@@ -621,6 +622,7 @@ def _expand(self, expand_input: ExpandInput, *, strict: bool) -> XComArg:
         task_id = get_unique_task_id(partial_kwargs.pop("task_id"), dag, task_group)
         if task_group:
             task_id = task_group.child_id(task_id)
+        validate_key(task_id)
 
         # Logic here should be kept in sync with BaseOperatorMeta.partial().
         if partial_kwargs.get("wait_for_downstream"):

task-sdk/src/airflow/sdk/bases/operator.py

@@ -345,6 +345,7 @@ def partial(
             task_group = task_group or TaskGroupContext.get_current(dag)
         if task_group:
             task_id = task_group.child_id(task_id)
+        validate_key(task_id)
 
         # Merge Dag and task group level defaults into user-supplied values.
         dag_default_args, partial_params = get_merged_defaults(

task-sdk/tests/task_sdk/definitions/decorators/test_mapped.py

@@ -17,9 +17,32 @@
 # under the License.
 from __future__ import annotations
 
+import pytest
+
 from airflow.sdk import DAG, TaskGroup
 
 
+def test_mapped_taskflow_task_id_is_validated():
+    def f(z):
+        pass
+
+    with DAG(dag_id="d", schedule=None) as dag:
+        with pytest.raises(ValueError, match="has to be made of alphanumeric"):
+            dag.task(task_id="bad*id")(f).expand(z=[])
+
+
+def test_mapped_taskflow_task_id_validated_after_unique_suffix():
+    # A 250-char id is valid the first time; the auto-added "__1" dedup suffix on the second
+    # use makes it 253 chars, so this only fails if the id is validated after deduplication.
+    def f(z):
+        pass
+
+    with DAG(dag_id="d", schedule=None) as dag:
+        dag.task(task_id="t" * 250)(f).expand(z=[])
+        with pytest.raises(ValueError, match="has to be less than 250 characters"):
+            dag.task(task_id="t" * 250)(f).expand(z=[])
+
+
 def test_mapped_task_group_id_prefix_task_id():
     def f(z):
         pass

task-sdk/tests/task_sdk/definitions/test_mappedoperator.py

@@ -75,6 +75,31 @@ def test_task_mapping_with_dag():
 # test_task_mapping_with_dag_and_list_of_pandas_dataframe
 
 
+@pytest.mark.parametrize(
+    ("bad_task_id", "match"),
+    [
+        ("bad*id", "has to be made of alphanumeric"),
+        ("with space", "has to be made of alphanumeric"),
+        ("a" * 251, "has to be less than 250 characters"),
+    ],
+)
+def test_mapped_task_id_is_validated(bad_task_id, match):
+    with DAG("test-dag"):
+        with pytest.raises(ValueError, match=match):
+            MockOperator.partial(task_id=bad_task_id).expand(arg2=["a"])
+
+
+def test_mapped_task_id_validated_after_group_prefix():
+    # A 50-char task id is valid on its own but is 251 chars once prefixed with the 200-char
+    # group id, so this only fails if the *final* prefixed id is what gets validated.
+    from airflow.sdk.definitions.taskgroup import TaskGroup
+
+    with DAG("test-dag"):
+        with TaskGroup("g" * 200):
+            with pytest.raises(ValueError, match="has to be less than 250 characters"):
+                MockOperator.partial(task_id="t" * 50).expand(arg2=["a"])
+
+
 def test_task_mapping_without_dag_context():
     with DAG("test-dag") as dag:
         task1 = BaseOperator(task_id="op1")