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)