From a7fb9b7b4c6f76c3a50f9c482bcc96ca437d0aa1 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Wed, 12 Aug 2026 17:05:54 -0700 Subject: [PATCH] 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> --- .../serialization/definitions/operatorlink.py | 48 ++++++--- .../definitions/test_operatorlink.py | 101 ++++++++++++++++++ .../src/airflow_shared/state/__init__.py | 9 ++ .../airflow/sdk/execution_time/task_runner.py | 5 + .../execution_time/test_task_runner.py | 29 +++++ 5 files changed, 180 insertions(+), 12 deletions(-) create mode 100644 airflow-core/tests/unit/serialization/definitions/test_operatorlink.py diff --git a/airflow-core/src/airflow/serialization/definitions/operatorlink.py b/airflow-core/src/airflow/serialization/definitions/operatorlink.py index b7a6eef11b84f..f462228dde53e 100644 --- a/airflow-core/src/airflow/serialization/definitions/operatorlink.py +++ b/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)) 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..3ca2a4181324a --- /dev/null +++ b/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" diff --git a/shared/state/src/airflow_shared/state/__init__.py b/shared/state/src/airflow_shared/state/__init__.py index e2c094d84818d..06b255a33b899 100644 --- a/shared/state/src/airflow_shared/state/__init__.py +++ b/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: """ 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..a7efb0f3f1d72 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/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", 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 cc9fb77e08921..62dd93d691881 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 @@ -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)