From 664061b33c1a231d1b2b5bea24195befc4cd0d49 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Tue, 11 Aug 2026 15:58:12 -0700 Subject: [PATCH 1/2] Return each attempt's own URL from operator extra links Operator extra links are cached as XCom rows under the link's xcom_key, which carries no attempt. XComOperatorLink.get_link receives a TaskInstanceKey and ignores its try_number, so once a task retries, asking for an earlier attempt's link returns the latest attempt's URL. #65661 made the endpoint resolve the requested attempt, which makes the wrong answer reachable from the UI. The worker now also caches the link under a key carrying the attempt that produced it, and the reader prefers that row, falling back to the bare key for links written before this change. Those rows are exempt from the clear a retry issues, since they describe attempts that already ran. The prefix is owned by Airflow rather than derived from the link's xcom_key, because a link may override that with any name, as the Databricks and Google links do. Closes: #71471 Signed-off-by: 1fanwang <1fannnw@gmail.com> --- airflow-core/newsfragments/71471.bugfix.rst | 1 + .../execution_api/routes/task_instances.py | 8 +++ .../serialization/definitions/operatorlink.py | 23 +++--- .../versions/head/test_task_instances.py | 46 ++++++++++++ .../definitions/test_operatorlink.py | 70 +++++++++++++++++++ .../src/airflow/sdk/bases/operatorlink.py | 19 +++++ .../airflow/sdk/execution_time/task_runner.py | 5 ++ 7 files changed, 163 insertions(+), 9 deletions(-) create mode 100644 airflow-core/newsfragments/71471.bugfix.rst create mode 100644 airflow-core/tests/unit/serialization/definitions/test_operatorlink.py diff --git a/airflow-core/newsfragments/71471.bugfix.rst b/airflow-core/newsfragments/71471.bugfix.rst new file mode 100644 index 0000000000000..90df564acceb3 --- /dev/null +++ b/airflow-core/newsfragments/71471.bugfix.rst @@ -0,0 +1 @@ +Return each attempt's own URL from operator extra links, instead of the latest attempt's. diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index 02d723a82185e..4ded029c440c6 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -94,6 +94,7 @@ from airflow.models.taskreschedule import TaskReschedule from airflow.models.trigger import Trigger, handle_event_submit from airflow.models.xcom import XComModel +from airflow.sdk.bases.operatorlink import ATTEMPT_LINK_XCOM_KEY_PREFIX from airflow.serialization.definitions.assets import SerializedAsset, SerializedAssetUniqueKey from airflow.state import get_state_backend from airflow.triggers.base import TriggerEvent @@ -296,6 +297,13 @@ def ti_run( if map_index is not None: xcom_query = xcom_query.where(XComModel.map_index == map_index) + # Each attempt's rendered operator links describe where that attempt ran, so + # they have to outlive it. The underscores are LIKE single-character wildcards + # and must be escaped, or unrelated keys are spared too. + xcom_query = xcom_query.where( + ~XComModel.key.like(ATTEMPT_LINK_XCOM_KEY_PREFIX.replace("_", r"\_") + "%", escape="\\") + ) + xcom_keys = list(session.scalars(xcom_query)) task_reschedule_count = ( session.scalar( diff --git a/airflow-core/src/airflow/serialization/definitions/operatorlink.py b/airflow-core/src/airflow/serialization/definitions/operatorlink.py index b7a6eef11b84f..6843d889b9001 100644 --- a/airflow-core/src/airflow/serialization/definitions/operatorlink.py +++ b/airflow-core/src/airflow/serialization/definitions/operatorlink.py @@ -24,6 +24,7 @@ import attrs from airflow.models.xcom import XComModel +from airflow.sdk.bases.operatorlink import attempt_link_xcom_key from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.session import create_session @@ -54,16 +55,20 @@ def get_link(self, operator: Operator, *, ti_key: TaskInstanceKey) -> str: self.log.info( "Attempting to retrieve link from XComs with key: %s for task id: %s", self.xcom_key, ti_key ) + keys = [attempt_link_xcom_key(self.xcom_key, ti_key.try_number), self.xcom_key] with create_session() as session: - result = session.execute( - XComModel.get_many( - key=self.xcom_key, - run_id=ti_key.run_id, - dag_ids=ti_key.dag_id, - task_ids=ti_key.task_id, - map_indexes=ti_key.map_index, - ).with_only_columns(XComModel.value) - ).first() + for key in keys: + result = session.execute( + XComModel.get_many( + key=key, + run_id=ti_key.run_id, + dag_ids=ti_key.dag_id, + task_ids=ti_key.task_id, + map_indexes=ti_key.map_index, + ).with_only_columns(XComModel.value) + ).first() + if result: + break if not result: self.log.debug( "No link with name: %s present in XCom as key: %s, returning empty link", diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index 80392fd17e206..e6cfd182ef566 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -53,8 +53,10 @@ from airflow.models.task_state_store import TaskStateStoreModel from airflow.models.taskinstance import TaskInstance from airflow.models.taskinstancehistory import TaskInstanceHistory +from airflow.models.xcom import XComModel from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk import Asset, TaskGroup, TriggerRule, task, task_group +from airflow.sdk.bases.operatorlink import attempt_link_xcom_key from airflow.state.metastore import MetastoreBackend from airflow.utils.state import DagRunState, State, TaskInstanceState, TerminalTIState @@ -886,6 +888,50 @@ def test_next_kwargs_determines_start_date_update(self, client, session, create_ ti = session.get(TaskInstance, ti.id) assert ti.start_date == expected_start_date + def test_ti_run_retains_per_attempt_operator_link_xcoms( + self, + client, + session, + create_task_instance, + ): + """A retry clears the task's XComs, but not each attempt's rendered links.""" + ti = create_task_instance( + task_id="test_ti_run_retains_per_attempt_operator_link_xcoms", + state=State.QUEUED, + session=session, + dag_id=str(uuid4()), + ) + retained = [ + attempt_link_xcom_key("_link_MyLink", 1), + attempt_link_xcom_key("databricks_job_run_link", 1), + ] + # "_" is a LIKE single-character wildcard, so these are spared too if it is unescaped. + cleared = ["return_value", "_link_MyLink", "Xlink_attemptX1_y", "my_link_attempt_1_thing"] + for key in retained + cleared: + XComModel.set( + key=key, + value="https://example.com", + dag_id=ti.dag_id, + task_id=ti.task_id, + run_id=ti.run_id, + session=session, + ) + session.commit() + + response = client.patch( + f"/execution/task-instances/{ti.id}/run", + json={ + "state": "running", + "hostname": "random-hostname", + "unixname": "random-unixname", + "pid": 100, + "start_date": DEFAULT_START_DATE.isoformat(), + }, + ) + + assert response.status_code == 200 + assert sorted(response.json()["xcom_keys_to_clear"]) == sorted(cleared) + def test_ti_run_resume_returns_original_start_date_in_context( self, client, diff --git a/airflow-core/tests/unit/serialization/definitions/test_operatorlink.py b/airflow-core/tests/unit/serialization/definitions/test_operatorlink.py new file mode 100644 index 0000000000000..2d724713afb8b --- /dev/null +++ b/airflow-core/tests/unit/serialization/definitions/test_operatorlink.py @@ -0,0 +1,70 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import pytest + +from airflow.models.xcom import XComModel +from airflow.sdk.bases.operatorlink import attempt_link_xcom_key +from airflow.serialization.definitions.operatorlink import XComOperatorLink + +pytestmark = pytest.mark.db_test + +XCOM_KEY = "_link_MyLink" + + +@pytest.fixture +def link(): + return XComOperatorLink(name="My Link", xcom_key=XCOM_KEY) + + +def _set(session, ti, key, value): + XComModel.set( + key=key, + value=value, + dag_id=ti.dag_id, + task_id=ti.task_id, + run_id=ti.run_id, + map_index=ti.map_index, + session=session, + ) + + +class TestXComOperatorLinkPerAttempt: + def test_returns_the_requested_attempts_link(self, session, create_task_instance, link): + ti = create_task_instance(task_id="test_link_per_attempt") + _set(session, ti, attempt_link_xcom_key(XCOM_KEY, 1), "https://logs/attempt-1") + _set(session, ti, attempt_link_xcom_key(XCOM_KEY, 2), "https://logs/attempt-2") + _set(session, ti, XCOM_KEY, "https://logs/attempt-2") + session.commit() + + assert link.get_link(ti.task, ti_key=ti.key._replace(try_number=1)) == "https://logs/attempt-1" + assert link.get_link(ti.task, ti_key=ti.key._replace(try_number=2)) == "https://logs/attempt-2" + + def test_falls_back_to_the_bare_key(self, session, create_task_instance, link): + """Links written before per-attempt rows existed only have the bare key.""" + ti = create_task_instance(task_id="test_link_fallback") + _set(session, ti, XCOM_KEY, "https://logs/only-one") + session.commit() + + assert link.get_link(ti.task, ti_key=ti.key._replace(try_number=1)) == "https://logs/only-one" + + def test_returns_empty_when_nothing_stored(self, session, create_task_instance, link): + ti = create_task_instance(task_id="test_link_missing") + session.commit() + + assert link.get_link(ti.task, ti_key=ti.key) == "" diff --git a/task-sdk/src/airflow/sdk/bases/operatorlink.py b/task-sdk/src/airflow/sdk/bases/operatorlink.py index ebc5a4dcd04db..6bb3ea9410c6f 100644 --- a/task-sdk/src/airflow/sdk/bases/operatorlink.py +++ b/task-sdk/src/airflow/sdk/bases/operatorlink.py @@ -27,6 +27,25 @@ from airflow.sdk.types import TaskInstanceKey +ATTEMPT_LINK_XCOM_KEY_PREFIX = "_link_attempt_" +"""Prefix marking the XCom rows that hold one attempt's rendered operator link. + +The prefix is owned by Airflow rather than derived from the link's own ``xcom_key``, +because a link is free to override that with any name it likes. +""" + + +def attempt_link_xcom_key(xcom_key: str, try_number: int) -> str: + """ + Return the XCom key holding ``xcom_key``'s link as rendered for ``try_number``. + + Example:: + + attempt_link_xcom_key("_link_MyLink", 2) == "_link_attempt_2__link_MyLink" + """ + return f"{ATTEMPT_LINK_XCOM_KEY_PREFIX}{try_number}_{xcom_key}" + + @attrs.define() class BaseOperatorLink(metaclass=ABCMeta): """Abstract base class that defines how we get an operator link.""" 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 1ce848128d663..ab0990bce4071 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -58,6 +58,7 @@ TIRunContext, ) from airflow.sdk.bases.operator import BaseOperator, ExecutorSafeguard +from airflow.sdk.bases.operatorlink import attempt_link_xcom_key from airflow.sdk.bases.xcom import BaseXCom from airflow.sdk.configuration import conf from airflow.sdk.definitions._internal.dag_parsing_context import _airflow_parsing_context_manager @@ -2329,6 +2330,10 @@ def finalize( link, xcom_key = oe.get_link(operator=task, ti_key=ti), oe.xcom_key # type: ignore[arg-type] log.debug("Setting xcom for operator extra link", link=link, xcom_key=xcom_key) _xcom_push_to_db(ti, key=xcom_key, value=link) + # The bare key holds the latest attempt and is overwritten by the next one, so + # also keep this attempt's link under its own key. Without it a retry leaves the + # UI resolving every earlier attempt's link to the last attempt's URL. + _xcom_push_to_db(ti, key=attempt_link_xcom_key(xcom_key, ti.try_number), value=link) except Exception: log.exception( "Failed to push an xcom for task operator extra link", From dc319919456bf3ac56aa97f5387ac2d6a2ed6647 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Wed, 12 Aug 2026 16:44:54 -0700 Subject: [PATCH 2/2] Drop the newsfragment while this is a draft Signed-off-by: 1fanwang <1fannnw@gmail.com> --- airflow-core/newsfragments/71471.bugfix.rst | 1 - 1 file changed, 1 deletion(-) delete mode 100644 airflow-core/newsfragments/71471.bugfix.rst diff --git a/airflow-core/newsfragments/71471.bugfix.rst b/airflow-core/newsfragments/71471.bugfix.rst deleted file mode 100644 index 90df564acceb3..0000000000000 --- a/airflow-core/newsfragments/71471.bugfix.rst +++ /dev/null @@ -1 +0,0 @@ -Return each attempt's own URL from operator extra links, instead of the latest attempt's.