From 581708c8de4482e5bc35bacfc240594758b9e98b Mon Sep 17 00:00:00 2001 From: Andrew Chang Date: Mon, 20 Jul 2026 22:31:09 +0800 Subject: [PATCH] Validate task_id for mapped tasks A regular task validates its task_id in BaseOperator.__init__, but a mapped task is built through a different path that skips __init__, so its task_id was never checked. An invalid id such as "bad*id" passed silently in .partial().expand() and in the taskflow @task mapped form, only to fail later where the id is used. Validate the id in both mapped paths so it is rejected at parse time like a regular task. --- task-sdk/src/airflow/sdk/bases/decorator.py | 2 ++ task-sdk/src/airflow/sdk/bases/operator.py | 1 + .../definitions/decorators/test_mapped.py | 23 +++++++++++++++++ .../definitions/test_mappedoperator.py | 25 +++++++++++++++++++ 4 files changed, 51 insertions(+) diff --git a/task-sdk/src/airflow/sdk/bases/decorator.py b/task-sdk/src/airflow/sdk/bases/decorator.py index e45778725cc16..15b94ee1e0b5c 100644 --- a/task-sdk/src/airflow/sdk/bases/decorator.py +++ b/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"): diff --git a/task-sdk/src/airflow/sdk/bases/operator.py b/task-sdk/src/airflow/sdk/bases/operator.py index 97da8869686cc..a18ba0d6f20d4 100644 --- a/task-sdk/src/airflow/sdk/bases/operator.py +++ b/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( diff --git a/task-sdk/tests/task_sdk/definitions/decorators/test_mapped.py b/task-sdk/tests/task_sdk/definitions/decorators/test_mapped.py index 14c3cd2fb1c27..083dfb1e8787c 100644 --- a/task-sdk/tests/task_sdk/definitions/decorators/test_mapped.py +++ b/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 diff --git a/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py b/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py index 4aadf923e9c85..b3429fdc0316d 100644 --- a/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py +++ b/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")