Return each attempt's own URL from operator extra links

An operator extra link is cached as a single XCom row under the link's xcom_key,
and XComOperatorLink.get_link ignores the try_number on the TaskInstanceKey it is
given. A task's XComs are also cleared before every attempt, so that row only ever
holds whichever attempt ran last. Asking for an earlier attempt's link therefore
returns the latest attempt's URL, and the original attempt's logs cannot be reached
from the UI at all.

The worker now also writes each attempt's rendered link to the task state store,
which a retry does not clear, and the reader prefers that row. The XCom row is left
in place and is still the answer for links written before this change.

Reading goes through get_state_backend(), the same resolver the execution API uses,
so a deployment pointing [state_store] backend elsewhere keeps working.

Closes: #71471
Signed-off-by: 1fanwang <1fannnw@gmail.com>

Author: 1fanwang

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

@@ -24,6 +24,9 @@
 import attrs
 
 from airflow.models.xcom import XComModel
+from airflow.sdk.bases.operatorlink import attempt_link_state_key
+from airflow._shared.state import TaskScope
+from airflow.state import get_state_backend
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.session import create_session
 
@@ -43,19 +46,28 @@ class XComOperatorLink(LoggingMixin):
     name: str
     xcom_key: str
 
-    def get_link(self, operator: Operator, *, ti_key: TaskInstanceKey) -> str:
-        """
-        Retrieve the link from the XComs.
+    def _stored_link(self, ti_key: TaskInstanceKey) -> str | None:
+        """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,102 @@
+# 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
+from airflow.models.xcom import XComModel
+from airflow.sdk.bases.operatorlink import attempt_link_state_key
+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(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_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(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="link_missing")
+        session.commit()
+
+        assert link.get_link(ti.task, 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(ti.task, ti_key=ti.key._replace(try_number=1)) == "https://logs/attempt-1"

task-sdk/src/airflow/sdk/bases/operatorlink.py

@@ -27,6 +27,15 @@
     from airflow.sdk.types import TaskInstanceKey
 
 
+ATTEMPT_LINK_STATE_KEY_PREFIX = "_link_attempt_"
+"""Prefix for the task-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}"
+
+
 @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_state_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 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.bases.operatorlink 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)