Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Licensed to the Apache Software Foundation (ASF) under one
2# or more contributor license agreements. See the NOTICE file
3# distributed with this work for additional information
4# regarding copyright ownership. The ASF licenses this file
5# to you under the Apache License, Version 2.0 (the
6# "License"); you may not use this file except in compliance
7# with the License. You may obtain a copy of the License at
8#
9# http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing,
12# software distributed under the License is distributed on an
13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14# KIND, either express or implied. See the License for the
15# specific language governing permissions and limitations
16# under the License.
17"""
18Airflow compatibility imports for seamless migration from Airflow 2 to Airflow 3.
20This module provides lazy imports that automatically try Airflow 3 paths first,
21then fall back to Airflow 2 paths, enabling code to work across both versions.
22"""
24from __future__ import annotations
26from typing import TYPE_CHECKING
28from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_0_PLUS
30if TYPE_CHECKING:
31 import airflow.sdk.io as io # noqa: F401
32 import airflow.sdk.timezone as timezone # noqa: F401
33 from airflow.models.xcom import XCOM_RETURN_KEY as XCOM_RETURN_KEY
34 from airflow.sdk import (
35 DAG as DAG,
36 Asset as Asset,
37 AssetAlias as AssetAlias,
38 AssetAll as AssetAll,
39 AssetAny as AssetAny,
40 BaseHook as BaseHook,
41 BaseNotifier as BaseNotifier,
42 BaseOperator as BaseOperator,
43 BaseOperatorLink as BaseOperatorLink,
44 BaseSensorOperator as BaseSensorOperator,
45 Connection as Connection,
46 Context as Context,
47 DagRunState as DagRunState,
48 EdgeModifier as EdgeModifier,
49 Label as Label,
50 Metadata as Metadata,
51 ObjectStoragePath as ObjectStoragePath,
52 Param as Param,
53 PokeReturnValue as PokeReturnValue,
54 TaskGroup as TaskGroup,
55 TaskInstanceState as TaskInstanceState,
56 TriggerRule as TriggerRule,
57 Variable as Variable,
58 WeightRule as WeightRule,
59 XComArg as XComArg,
60 chain as chain,
61 chain_linear as chain_linear,
62 cross_downstream as cross_downstream,
63 dag as dag,
64 get_current_context as get_current_context,
65 get_parsing_context as get_parsing_context,
66 setup as setup,
67 task as task,
68 task_group as task_group,
69 teardown as teardown,
70 )
71 from airflow.sdk.bases.decorator import (
72 DecoratedMappedOperator as DecoratedMappedOperator,
73 DecoratedOperator as DecoratedOperator,
74 TaskDecorator as TaskDecorator,
75 get_unique_task_id as get_unique_task_id,
76 task_decorator_factory as task_decorator_factory,
77 )
78 from airflow.sdk.bases.sensor import poke_mode_only as poke_mode_only
79 from airflow.sdk.definitions.context import context_merge as context_merge
80 from airflow.sdk.definitions.mappedoperator import MappedOperator as MappedOperator
81 from airflow.sdk.definitions.template import literal as literal
82 from airflow.sdk.exceptions import (
83 AirflowException as AirflowException,
84 AirflowFailException as AirflowFailException,
85 AirflowNotFoundException as AirflowNotFoundException,
86 AirflowSensorTimeout as AirflowSensorTimeout,
87 AirflowSkipException as AirflowSkipException,
88 AirflowTaskTimeout as AirflowTaskTimeout,
89 ParamValidationError as ParamValidationError,
90 TaskDeferred as TaskDeferred,
91 XComNotFound as XComNotFound,
92 )
93 from airflow.sdk.log import redact as redact
94 from airflow.sdk.observability.stats import Stats as Stats
96 # Airflow 3-only exceptions (conditionally imported)
97 if AIRFLOW_V_3_0_PLUS:
98 from airflow.sdk.exceptions import (
99 DagRunTriggerException as DagRunTriggerException,
100 DownstreamTasksSkipped as DownstreamTasksSkipped,
101 )
102 from airflow.sdk.execution_time.context import (
103 AIRFLOW_VAR_NAME_FORMAT_MAPPING as AIRFLOW_VAR_NAME_FORMAT_MAPPING,
104 context_to_airflow_vars as context_to_airflow_vars,
105 )
106 from airflow.sdk.execution_time.timeout import timeout as timeout
107 from airflow.sdk.execution_time.xcom import XCom as XCom
110from airflow.providers.common.compat._compat_utils import create_module_getattr
112# Rename map for classes that changed names between Airflow 2.x and 3.x
113# Format: new_name -> (new_path, old_path, old_name)
114_RENAME_MAP: dict[str, tuple[str, str, str]] = {
115 # Assets: Dataset -> Asset rename in Airflow 3.0
116 "Asset": ("airflow.sdk", "airflow.datasets", "Dataset"),
117 "AssetAlias": ("airflow.sdk", "airflow.datasets", "DatasetAlias"),
118 "AssetAll": ("airflow.sdk", "airflow.datasets", "DatasetAll"),
119 "AssetAny": ("airflow.sdk", "airflow.datasets", "DatasetAny"),
120}
122# Import map for classes/functions/constants
123# Format: class_name -> module_path(s)
124# - str: single module path (no fallback)
125# - tuple[str, ...]: multiple module paths (try in order, newest first)
126_IMPORT_MAP: dict[str, str | tuple[str, ...]] = {
127 # ============================================================================
128 # Hooks
129 # ============================================================================
130 "BaseHook": ("airflow.sdk", "airflow.hooks.base"),
131 # ============================================================================
132 # Sensors
133 # ============================================================================
134 "BaseSensorOperator": ("airflow.sdk", "airflow.sensors.base"),
135 "PokeReturnValue": ("airflow.sdk", "airflow.sensors.base"),
136 "poke_mode_only": ("airflow.sdk.bases.sensor", "airflow.sensors.base"),
137 # ============================================================================
138 # Operators
139 # ============================================================================
140 "BaseOperator": ("airflow.sdk", "airflow.models.baseoperator"),
141 # ============================================================================
142 # Decorators
143 # ============================================================================
144 "task": ("airflow.sdk", "airflow.decorators"),
145 "dag": ("airflow.sdk", "airflow.decorators"),
146 "task_group": ("airflow.sdk", "airflow.decorators"),
147 "setup": ("airflow.sdk", "airflow.decorators"),
148 "teardown": ("airflow.sdk", "airflow.decorators"),
149 "TaskDecorator": ("airflow.sdk.bases.decorator", "airflow.decorators"),
150 "task_decorator_factory": ("airflow.sdk.bases.decorator", "airflow.decorators.base"),
151 "get_unique_task_id": ("airflow.sdk.bases.decorator", "airflow.decorators.base"),
152 # ============================================================================
153 # Models
154 # ============================================================================
155 "Connection": ("airflow.sdk", "airflow.models.connection"),
156 "Variable": ("airflow.sdk", "airflow.models.variable"),
157 "XCom": ("airflow.sdk.execution_time.xcom", "airflow.models.xcom"),
158 "DAG": ("airflow.sdk", "airflow.models.dag"),
159 "Param": ("airflow.sdk", "airflow.models.param"),
160 "XComArg": ("airflow.sdk", "airflow.models.xcom_arg"),
161 "DecoratedOperator": ("airflow.sdk.bases.decorator", "airflow.decorators.base"),
162 "DecoratedMappedOperator": ("airflow.sdk.bases.decorator", "airflow.decorators.base"),
163 "MappedOperator": ("airflow.sdk.definitions.mappedoperator", "airflow.models.mappedoperator"),
164 # ============================================================================
165 # Assets (Dataset → Asset rename in Airflow 3.0)
166 # ============================================================================
167 # Note: Asset, AssetAlias, AssetAll, AssetAny are handled by _RENAME_MAP
168 # Metadata moved from airflow.datasets.metadata (2.x) to airflow.sdk (3.x)
169 "Metadata": ("airflow.sdk", "airflow.datasets.metadata"),
170 # ============================================================================
171 # Notifiers
172 # ============================================================================
173 "BaseNotifier": ("airflow.sdk", "airflow.notifications.basenotifier"),
174 # ============================================================================
175 # Operator Links & Task Groups
176 # ============================================================================
177 "BaseOperatorLink": ("airflow.sdk", "airflow.models.baseoperatorlink"),
178 "TaskGroup": ("airflow.sdk", "airflow.utils.task_group"),
179 # ============================================================================
180 # Operator Utilities (chain, cross_downstream, etc.)
181 # ============================================================================
182 "chain": ("airflow.sdk", "airflow.models.baseoperator"),
183 "chain_linear": ("airflow.sdk", "airflow.models.baseoperator"),
184 "cross_downstream": ("airflow.sdk", "airflow.models.baseoperator"),
185 # ============================================================================
186 # Edge Modifiers & Labels
187 # ============================================================================
188 "EdgeModifier": ("airflow.sdk", "airflow.utils.edgemodifier"),
189 "Label": ("airflow.sdk", "airflow.utils.edgemodifier"),
190 # ============================================================================
191 # State Enums
192 # ============================================================================
193 "DagRunState": ("airflow.sdk", "airflow.utils.state"),
194 "TaskInstanceState": ("airflow.sdk", "airflow.utils.state"),
195 "TriggerRule": ("airflow.sdk", "airflow.utils.trigger_rule"),
196 "WeightRule": ("airflow.sdk", "airflow.utils.weight_rule"),
197 # ============================================================================
198 # IO & Storage
199 # ============================================================================
200 "ObjectStoragePath": ("airflow.sdk", "airflow.io.path"),
201 # ============================================================================
202 # Template Utilities
203 # ============================================================================
204 "literal": ("airflow.sdk.definitions.template", "airflow.utils.template"),
205 # ============================================================================
206 # Context & Utilities
207 # ============================================================================
208 "Context": ("airflow.sdk", "airflow.utils.context"),
209 "context_merge": ("airflow.sdk.definitions.context", "airflow.utils.context"),
210 "context_to_airflow_vars": ("airflow.sdk.execution_time.context", "airflow.utils.operator_helpers"),
211 "AIRFLOW_VAR_NAME_FORMAT_MAPPING": (
212 "airflow.sdk.execution_time.context",
213 "airflow.utils.operator_helpers",
214 ),
215 "get_current_context": ("airflow.sdk", "airflow.operators.python"),
216 "get_parsing_context": ("airflow.sdk", "airflow.utils.dag_parsing_context"),
217 # ============================================================================
218 # Timeout Utilities
219 # ============================================================================
220 "timeout": ("airflow.sdk.execution_time.timeout", "airflow.utils.timeout"),
221 # ============================================================================
222 # XCom & Task Communication
223 # ============================================================================
224 "XCOM_RETURN_KEY": "airflow.models.xcom",
225 # ============================================================================
226 # Exceptions (deprecated in airflow.exceptions, prefer SDK)
227 # ============================================================================
228 # Note: AirflowException and AirflowNotFoundException are not deprecated, but exposing them
229 # here keeps provider imports consistent across Airflow 2 and 3.
230 "AirflowException": ("airflow.sdk.exceptions", "airflow.exceptions"),
231 "AirflowFailException": ("airflow.sdk.exceptions", "airflow.exceptions"),
232 "AirflowNotFoundException": ("airflow.sdk.exceptions", "airflow.exceptions"),
233 "AirflowSkipException": ("airflow.sdk.exceptions", "airflow.exceptions"),
234 "AirflowTaskTimeout": ("airflow.sdk.exceptions", "airflow.exceptions"),
235 "AirflowSensorTimeout": ("airflow.sdk.exceptions", "airflow.exceptions"),
236 "ParamValidationError": ("airflow.sdk.exceptions", "airflow.exceptions"),
237 "TaskDeferred": ("airflow.sdk.exceptions", "airflow.exceptions"),
238 "XComNotFound": ("airflow.sdk.exceptions", "airflow.exceptions"),
239 # ============================================================================
240 # Observability
241 # ============================================================================
242 "Stats": ("airflow.sdk.observability.stats", "airflow.stats"),
243 # ============================================================================
244 # Secrets Masking
245 # ============================================================================
246 "redact": (
247 "airflow.sdk.log",
248 "airflow.sdk._shared.secrets_masker",
249 "airflow.sdk.execution_time.secrets_masker",
250 "airflow.utils.log.secrets_masker",
251 ),
252}
254# Airflow 3-only exceptions (not available in Airflow 2)
255_AIRFLOW_3_ONLY_EXCEPTIONS: dict[str, tuple[str, ...]] = {
256 "DownstreamTasksSkipped": ("airflow.sdk.exceptions", "airflow.exceptions"),
257 "DagRunTriggerException": ("airflow.sdk.exceptions", "airflow.exceptions"),
258}
260# Add Airflow 3-only exceptions to _IMPORT_MAP if running Airflow 3+
261if AIRFLOW_V_3_0_PLUS:
262 _IMPORT_MAP.update(_AIRFLOW_3_ONLY_EXCEPTIONS)
264# Module map: module_name -> module_path(s)
265# For entire modules that have been moved (e.g., timezone)
266# Usage: from airflow.providers.common.compat.lazy_compat import timezone
267_MODULE_MAP: dict[str, str | tuple[str, ...]] = {
268 "timezone": ("airflow.sdk.timezone", "airflow.utils.timezone"),
269 "io": ("airflow.sdk.io", "airflow.io"),
270}
272# Use the shared utility to create __getattr__
273__getattr__ = create_module_getattr(
274 import_map=_IMPORT_MAP,
275 module_map=_MODULE_MAP,
276 rename_map=_RENAME_MAP,
277)
279__all__ = list(_RENAME_MAP.keys()) + list(_IMPORT_MAP.keys()) + list(_MODULE_MAP.keys())