diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index 02d723a82185e..4ded029c440c6 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/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( diff --git a/airflow-core/src/airflow/serialization/definitions/operatorlink.py b/airflow-core/src/airflow/serialization/definitions/operatorlink.py index b7a6eef11b84f..6843d889b9001 100644 --- a/airflow-core/src/airflow/serialization/definitions/operatorlink.py +++ b/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", diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index 80392fd17e206..e6cfd182ef566 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/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, 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..2d724713afb8b --- /dev/null +++ b/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) == "" diff --git a/task-sdk/src/airflow/sdk/bases/operatorlink.py b/task-sdk/src/airflow/sdk/bases/operatorlink.py index ebc5a4dcd04db..6bb3ea9410c6f 100644 --- a/task-sdk/src/airflow/sdk/bases/operatorlink.py +++ b/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.""" 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..ab0990bce4071 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/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",