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..88fb19293d0c1 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 @@ -2334,8 +2335,8 @@ 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() - rendered_map_index = jinja_env.from_string(template).render(context) + 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 2cd63cce594c7..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,6 +2536,45 @@ def test_function(): assert ti.rendered_map_index == "Hello! test_run" +@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_numeric_map_index_template", + schedule=None, + render_template_as_native_obj=render_template_as_native_obj, + ): + task = PythonOperator( + task_id="test_task", + python_callable=test_function, + 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 == expected + + 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."""