diff --git a/airflow-core/src/airflow/utils/__init__.py b/airflow-core/src/airflow/utils/__init__.py index 574fbb642f35d..39a213c77c64f 100644 --- a/airflow-core/src/airflow/utils/__init__.py +++ b/airflow-core/src/airflow/utils/__init__.py @@ -43,5 +43,8 @@ def __getattr__(name: str): "BaseSetupTeardownContext": "airflow.sdk.definitions._internal.setup_teardown.BaseSetupTeardownContext", "SetupTeardownContext": "airflow.sdk.definitions._internal.setup_teardown.SetupTeardownContext", }, + "xcom": { + "XCOM_RETURN_KEY": "airflow.models.xcom.XCOM_RETURN_KEY", + }, } add_deprecated_classes(__deprecated_classes, __name__) diff --git a/airflow-core/src/airflow/utils/xcom.py b/airflow-core/src/airflow/utils/xcom.py deleted file mode 100644 index f65be31a834f1..0000000000000 --- a/airflow-core/src/airflow/utils/xcom.py +++ /dev/null @@ -1,24 +0,0 @@ -# -# 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. - -# MAX XCOM Size is 48KB -# https://github.com/apache/airflow/pull/1618#discussion_r68249677 -from __future__ import annotations - -# TODO: Remove this once all the providers have been moved to not use this import -XCOM_RETURN_KEY = "return_value" diff --git a/devel-common/src/tests_common/test_utils/compat.py b/devel-common/src/tests_common/test_utils/compat.py index e5a4638f95c30..383fcce62e0c5 100644 --- a/devel-common/src/tests_common/test_utils/compat.py +++ b/devel-common/src/tests_common/test_utils/compat.py @@ -57,6 +57,12 @@ from airflow.sensors.date_time import DateTimeSensor # type: ignore[no-redef] from airflow.utils.python_virtualenv import write_python_script # type: ignore[no-redef] +try: + from airflow.models.xcom import XCOM_RETURN_KEY +except ImportError: + # Compatibility for Airflow < 3.1 + from airflow.utils.xcom import XCOM_RETURN_KEY # type: ignore[no-redef,attr-defined] + if TYPE_CHECKING: from airflow.models import Connection diff --git a/devel-common/src/tests_common/test_utils/mock_context.py b/devel-common/src/tests_common/test_utils/mock_context.py index 200d8a583ab0b..2a22cd223a118 100644 --- a/devel-common/src/tests_common/test_utils/mock_context.py +++ b/devel-common/src/tests_common/test_utils/mock_context.py @@ -31,7 +31,8 @@ def mock_context(task) -> Context: from airflow.models import TaskInstance from airflow.utils.session import NEW_SESSION from airflow.utils.state import TaskInstanceState - from airflow.utils.xcom import XCOM_RETURN_KEY + + from tests_common.test_utils.compat import XCOM_RETURN_KEY values: dict[str, Any] = {} diff --git a/devel-common/src/tests_common/test_utils/version_compat.py b/devel-common/src/tests_common/test_utils/version_compat.py index 95a8c5a41a79d..55c045cd9491a 100644 --- a/devel-common/src/tests_common/test_utils/version_compat.py +++ b/devel-common/src/tests_common/test_utils/version_compat.py @@ -41,7 +41,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: XCOM_RETURN_KEY = BaseXCom.XCOM_RETURN_KEY else: - from airflow.utils.xcom import XCOM_RETURN_KEY + from airflow.utils.xcom import XCOM_RETURN_KEY # type: ignore[no-redef] def get_sqlalchemy_version_tuple() -> tuple[int, int, int]: diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/version_compat.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/version_compat.py index 2a974225f0f96..e14332d52f29b 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/version_compat.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/version_compat.py @@ -41,7 +41,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: else: from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] from airflow.models import BaseOperator - from airflow.utils.xcom import XCOM_RETURN_KEY + from airflow.utils.xcom import XCOM_RETURN_KEY # type: ignore[no-redef] if AIRFLOW_V_3_0_PLUS: from airflow.sdk import BaseSensorOperator diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/version_compat.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/version_compat.py index 1570bf8bffe1a..37bb81d3b5419 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/version_compat.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/version_compat.py @@ -52,7 +52,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: from airflow.sdk import BaseHook else: from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] - from airflow.utils.xcom import XCOM_RETURN_KEY + from airflow.utils.xcom import XCOM_RETURN_KEY # type: ignore[no-redef] __all__ = [ "AIRFLOW_V_3_0_PLUS",