Return each attempt's own URL from operator extra links

airflow-core/src/airflow/serialization/definitions/operatorlink.py

@@ -23,7 +23,9 @@
 
 import attrs
 
+from airflow._shared.state import TaskScope, attempt_link_state_key
 from airflow.models.xcom import XComModel
+from airflow.state import get_state_backend
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.session import create_session
 
@@ -43,19 +45,29 @@ class XComOperatorLink(LoggingMixin):
     name: str
     xcom_key: str
 
-    def get_link(self, operator: Operator, *, ti_key: TaskInstanceKey) -> str:
+    def _stored_link(self, ti_key: TaskInstanceKey) -> str | None:
         """
-        Retrieve the link from the XComs.
+        Return the stored link for ``ti_key``'s attempt, or None.
 
-        :param operator: The Airflow operator object this link is associated to.
-        :param ti_key: TaskInstance ID to return link for.
-        :return: link to external system, but by pulling it from XComs
+        The state store is read first because it is the only place an earlier attempt's link
+        survives: a task's XComs are cleared before every attempt, so the XCom row holds
+        whichever attempt ran last. That row is still the answer for links written before
+        per-attempt rows existed.
         """
-        self.log.info(
-            "Attempting to retrieve link from XComs with key: %s for task id: %s", self.xcom_key, ti_key
+        scope = TaskScope(
+            dag_id=ti_key.dag_id,
+            run_id=ti_key.run_id,
+            task_id=ti_key.task_id,
+            map_index=ti_key.map_index,
         )
         with create_session() as session:
-            result = session.execute(
+            stored = get_state_backend().get(
+                scope, attempt_link_state_key(self.xcom_key, ti_key.try_number), session=session
+            )
+            if stored is not None:
+                return stored
+
+            row = session.execute(
                 XComModel.get_many(
                     key=self.xcom_key,
                     run_id=ti_key.run_id,
@@ -64,9 +76,21 @@ def get_link(self, operator: Operator, *, ti_key: TaskInstanceKey) -> str:
                     map_indexes=ti_key.map_index,
                 ).with_only_columns(XComModel.value)
             ).first()
-        if not result:
+        return row.value if row else None
+
+    def get_link(self, operator: Operator, *, ti_key: TaskInstanceKey) -> str:
+        """
+        Retrieve the link from the XComs.
+
+        :param operator: The Airflow operator object this link is associated to.
+        :param ti_key: TaskInstance ID to return link for.
+        :return: link to external system, but by pulling it from XComs
+        """
+        self.log.info("Attempting to retrieve link with key: %s for task id: %s", self.xcom_key, ti_key)
+        raw_value = self._stored_link(ti_key)
+        if raw_value is None:
             self.log.debug(
-                "No link with name: %s present in XCom as key: %s, returning empty link",
+                "No link with name: %s present for key: %s, returning empty link",
                 self.name,
                 self.xcom_key,
             )
@@ -78,10 +102,10 @@ def get_link(self, operator: Operator, *, ti_key: TaskInstanceKey) -> str:
         )
 
         try:
-            parsed_value = json.loads(result.value)
+            parsed_value = json.loads(raw_value)
         except (ValueError, TypeError):
             # Handling for cases when types do not need to be deserialized (e.g. when value is a simple string link)
-            parsed_value = result.value
+            parsed_value = raw_value
 
         try:
             return str(stringify_xcom(parsed_value))

airflow-core/tests/unit/serialization/definitions/test_operatorlink.py

@@ -0,0 +1,101 @@
+# 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 json
+
+import pytest
+
+from airflow._shared.state import TaskScope, attempt_link_state_key
+from airflow.models.xcom import XComModel
+from airflow.serialization.definitions.operatorlink import XComOperatorLink
+from airflow.state import get_state_backend
+
+pytestmark = pytest.mark.db_test
+
+XCOM_KEY = "_link_MyLink"
+
+
+@pytest.fixture
+def link():
+    return XComOperatorLink(name="My Link", xcom_key=XCOM_KEY)
+
+
+@pytest.fixture
+def store_link(session):
+    def write(ti, try_number, value):
+        get_state_backend().set(
+            TaskScope(dag_id=ti.dag_id, run_id=ti.run_id, task_id=ti.task_id, map_index=ti.map_index),
+            attempt_link_state_key(XCOM_KEY, try_number),
+            json.dumps(value),
+            session=session,
+        )
+
+    return write
+
+
+@pytest.fixture
+def xcom_link(session):
+    def write(ti, value):
+        XComModel.set(
+            key=XCOM_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,
+        )
+
+    return write
+
+
+class TestXComOperatorLinkPerAttempt:
+    def test_returns_the_requested_attempts_link(
+        self, session, create_task_instance, link, store_link, xcom_link
+    ):
+        ti = create_task_instance(task_id="link_per_attempt")
+        store_link(ti, 1, "https://logs/attempt-1")
+        store_link(ti, 2, "https://logs/attempt-2")
+        xcom_link(ti, "https://logs/attempt-2")
+        session.commit()
+
+        assert link.get_link(None, ti_key=ti.key._replace(try_number=1)) == "https://logs/attempt-1"
+        assert link.get_link(None, ti_key=ti.key._replace(try_number=2)) == "https://logs/attempt-2"
+
+    def test_falls_back_to_xcom(self, session, create_task_instance, link, xcom_link):
+        """Links written before per-attempt rows existed only have the XCom row."""
+        ti = create_task_instance(task_id="link_fallback")
+        xcom_link(ti, "https://logs/only-one")
+        session.commit()
+
+        assert link.get_link(None, 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="link_missing")
+        session.commit()
+
+        assert link.get_link(None, ti_key=ti.key) == ""
+
+    def test_state_store_wins_over_xcom(self, session, create_task_instance, link, store_link, xcom_link):
+        """The XCom row is the latest attempt, so it must not answer for an earlier one."""
+        ti = create_task_instance(task_id="link_precedence")
+        store_link(ti, 1, "https://logs/attempt-1")
+        xcom_link(ti, "https://logs/attempt-2")
+        session.commit()
+
+        assert link.get_link(None, ti_key=ti.key._replace(try_number=1)) == "https://logs/attempt-1"

shared/state/src/airflow_shared/state/__init__.py

@@ -30,6 +30,15 @@
     from sqlalchemy.orm import Session
 
 
+ATTEMPT_LINK_STATE_KEY_PREFIX = "_link_attempt_"
+"""Prefix for the state-store keys holding one attempt's rendered operator link."""
+
+
+def attempt_link_state_key(xcom_key: str, try_number: int) -> str:
+    """Return the state-store key holding ``xcom_key``'s link as rendered for ``try_number``."""
+    return f"{ATTEMPT_LINK_STATE_KEY_PREFIX}{try_number}_{xcom_key}"
+
+
 @dataclass(frozen=True)
 class TaskScope:
     """

task-sdk/src/airflow/sdk/execution_time/task_runner.py

@@ -47,6 +47,7 @@
 from airflow.sdk._shared.observability.metrics import stats
 from airflow.sdk._shared.observability.metrics.stats import build_dag_metric_tags
 from airflow.sdk._shared.observability.traces import get_task_span_detail_level
+from airflow.sdk._shared.state import attempt_link_state_key
 from airflow.sdk._shared.template_rendering import truncate_rendered_value
 from airflow.sdk.api.client import get_hostname, getuser
 from airflow.sdk.api.datamodels._generated import (
@@ -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 task's XComs are cleared before the next attempt, so this attempt's link
+            # goes to the state store, which is not.
+            if (store := context.get("task_state_store")) is not None:
+                store.set(attempt_link_state_key(xcom_key, ti.try_number), link)
         except Exception:
             log.exception(
                 "Failed to push an xcom for task operator extra link",

task-sdk/tests/task_sdk/execution_time/test_task_runner.py

@@ -1196,6 +1196,35 @@ def execute(self, context):
     assert counted.count("operator_failures") == 1
 
 
+def test_finalize_stores_the_operator_link_per_attempt(
+    mocked_parse, create_runtime_ti, mock_supervisor_comms
+):
+    """A retry clears the task's XComs, so each attempt's link also goes to the state store."""
+    from airflow.sdk._shared.state import attempt_link_state_key
+
+    class MyLink(BaseOperatorLink):
+        name = "My Link"
+
+        def get_link(self, operator, *, ti_key):
+            return f"https://logs/attempt-{ti_key.try_number}"
+
+    class CustomOperator(BaseOperator):
+        operator_extra_links = (MyLink(),)
+
+        def execute(self, context):
+            return None
+
+    ti = create_runtime_ti(task=CustomOperator(task_id="link_per_attempt"))
+    ti.try_number = 2
+    context = ti.get_template_context()
+    store = mock.MagicMock()
+    context["task_state_store"] = store
+
+    finalize(ti, context=context, log=mock.MagicMock(), state=TaskInstanceState.SUCCESS)
+
+    store.set.assert_called_once_with(attempt_link_state_key("_link_MyLink", 2), "https://logs/attempt-2")
+
+
 def test_run_downstream_skipped(mocked_parse, create_runtime_ti, mock_supervisor_comms, listener_manager):
     listener = TestTaskRunnerCallsListeners.CustomListener()
     listener_manager(listener)