Operator extra links per attempt (draft A): a row per attempt in XCom

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(

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",

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,

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) == ""

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."""

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",