From 6da526b74e8dc83332fc954297a10b2ce62ff1db Mon Sep 17 00:00:00 2001 From: Prajwal7842 Date: Fri, 14 Aug 2026 12:15:13 +0530 Subject: [PATCH 1/2] Explicitly convert map_index to str --- .../airflow/sdk/execution_time/task_runner.py | 3 ++- .../execution_time/test_task_runner.py | 24 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index b90ec01e21e99..c4ea021ae28e8 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -72,6 +72,7 @@ AssetUniqueKey, AssetUriRef, ) +from airflow.sdk.definitions.context import render_template_to_string from airflow.sdk.definitions.mappedoperator import MappedOperator from airflow.sdk.definitions.param import process_params from airflow.sdk.definitions.xcom_arg import PlainXComArg @@ -2335,7 +2336,7 @@ def _render_map_index(context: Context, ti: RuntimeTaskInstance, log: Logger) -> return None log.debug("Rendering map_index_template", template_length=len(template)) jinja_env = ti.task.dag.get_template_env() - rendered_map_index = jinja_env.from_string(template).render(context) + rendered_map_index = render_template_to_string(jinja_env.from_string(template), context) log.debug("Map index rendered", length=len(rendered_map_index)) return rendered_map_index diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 2cd63cce594c7..4929d8f3b23d1 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -2536,6 +2536,30 @@ def test_function(): assert ti.rendered_map_index == "Hello! test_run" +def test_execute_success_task_with_numeric_rendered_map_index(create_runtime_ti, mock_supervisor_comms): + """Test that a numeric map index remains a string with native rendering enabled.""" + + def test_function(): + return "test function" + + with DAG( + dag_id="dag_with_native_map_index_template", + schedule=None, + render_template_as_native_obj=True, + ): + task = PythonOperator( + task_id="test_task", + python_callable=test_function, + map_index_template="123", + ) + + ti = create_runtime_ti(task) + + run(ti, ti.get_template_context(), log=mock.MagicMock()) + + assert ti.rendered_map_index == "123" + + def test_execute_failed_task_with_rendered_map_index(create_runtime_ti, mock_supervisor_comms): """Test that the map index is rendered in the task context.""" From 1f4d4c8483e0342045fc7e0f1c7ee2ec0d787487 Mon Sep 17 00:00:00 2001 From: Prajwal7842 Date: Sun, 27 Sep 2026 21:29:11 +0530 Subject: [PATCH 2/2] Address comments --- .../airflow/sdk/execution_time/task_runner.py | 2 +- .../execution_time/test_task_runner.py | 27 ++++++++++++++----- 2 files changed, 22 insertions(+), 7 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index c4ea021ae28e8..88fb19293d0c1 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -2335,7 +2335,7 @@ def _render_map_index(context: Context, ti: RuntimeTaskInstance, log: Logger) -> if (template := context.get("map_index_template")) is None: return None log.debug("Rendering map_index_template", template_length=len(template)) - jinja_env = ti.task.dag.get_template_env() + jinja_env = ti.task.dag.get_template_env(force_sandboxed=True) rendered_map_index = render_template_to_string(jinja_env.from_string(template), context) log.debug("Map index rendered", length=len(rendered_map_index)) return rendered_map_index diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 4929d8f3b23d1..8988428222178 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -2536,28 +2536,43 @@ def test_function(): assert ti.rendered_map_index == "Hello! test_run" -def test_execute_success_task_with_numeric_rendered_map_index(create_runtime_ti, mock_supervisor_comms): - """Test that a numeric map index remains a string with native rendering enabled.""" +@pytest.mark.parametrize( + ("map_index_template", "expected"), + [ + ("{{ ti.try_number }}", "1"), + ("index-{{ ti.try_number }}", "index-1"), + ("{{ '3.10' }}", "3.10"), + ], +) +@pytest.mark.parametrize("render_template_as_native_obj", [False, True]) +def test_execute_success_task_with_numeric_rendered_map_index( + create_runtime_ti, + mock_supervisor_comms, + map_index_template, + expected, + render_template_as_native_obj, +): + """Test that numeric-looking map indices remain strings in both rendering modes.""" def test_function(): return "test function" with DAG( - dag_id="dag_with_native_map_index_template", + dag_id="dag_with_numeric_map_index_template", schedule=None, - render_template_as_native_obj=True, + render_template_as_native_obj=render_template_as_native_obj, ): task = PythonOperator( task_id="test_task", python_callable=test_function, - map_index_template="123", + map_index_template=map_index_template, ) ti = create_runtime_ti(task) run(ti, ti.get_template_context(), log=mock.MagicMock()) - assert ti.rendered_map_index == "123" + assert ti.rendered_map_index == expected def test_execute_failed_task_with_rendered_map_index(create_runtime_ti, mock_supervisor_comms):