Skip to content

Commit c54c1a9

Browse files
kaxilkaxil
authored
Add comprehensive compatibility imports for Airflow 2 to 3 migration (apache#56790)
1 parent a7f5337 commit c54c1a9

File tree

9 files changed

+1046
-51
lines changed

9 files changed

+1046
-51
lines changed

‎.pre-commit-config.yaml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -465,6 +465,12 @@ repos:
465465
entry: ./scripts/ci/prek/check_airflow_imports.py
466466
--pattern '^openlineage\.client\.(facet|run)'
467467
--message "You should import from `airflow.providers.common.compat.openlineage.facet` instead."
468+
- id: check-common-compat-lazy-imports-in-sync
469+
name: Check common.compat lazy_compat.pyi is in sync
470+
language: python
471+
files: ^providers/common/compat/src/airflow/providers/common/compat/lazy_compat\.(py|pyi)$
472+
pass_filenames: false
473+
entry: ./scripts/ci/prek/check_common_compat_lazy_imports.py
468474
- id: check-airflow-providers-bug-report-template
469475
name: Sort airflow-bug-report provider list
470476
language: python

‎providers/common/compat/pyproject.toml‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,6 @@ dev = [
7676
"apache-airflow-task-sdk",
7777
"apache-airflow-devel-common",
7878
"apache-airflow-providers-openlineage",
79-
"apache-airflow-providers-standard",
8079
# Additional devel dependencies (do not remove this line and add extra development dependencies)
8180
]
8281

‎providers/common/compat/src/airflow/providers/common/compat/lazy_compat.py‎

Lines changed: 306 additions & 0 deletions
Large diffs are not rendered by default.

‎providers/common/compat/src/airflow/providers/common/compat/lazy_compat.pyi‎

Lines changed: 230 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,230 @@
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+
"""
18+
Type stubs for IDE autocomplete - always uses Airflow 3 paths.
19+
20+
This file is auto-generated from lazy_compat.py.
21+
- run scripts/ci/prek/check_common_compat_lazy_imports.py --generate instead.
22+
"""
23+
24+
import airflow.sdk.io as io
25+
import airflow.sdk.timezone as timezone
26+
from airflow.exceptions import (
27+
AirflowBadRequest as AirflowBadRequest,
28+
AirflowConfigException as AirflowConfigException,
29+
AirflowException as AirflowException,
30+
AirflowFailException as AirflowFailException,
31+
AirflowNotFoundException as AirflowNotFoundException,
32+
AirflowSensorTimeout as AirflowSensorTimeout,
33+
AirflowSkipException as AirflowSkipException,
34+
AirflowTaskTerminated as AirflowTaskTerminated,
35+
AirflowTaskTimeout as AirflowTaskTimeout,
36+
)
37+
from airflow.models.dagrun import DagRun as DagRun
38+
from airflow.models.mappedoperator import MappedOperator as MappedOperator
39+
from airflow.models.taskinstance import TaskInstance as TaskInstance
40+
from airflow.models.xcom import XCOM_RETURN_KEY as XCOM_RETURN_KEY
41+
from airflow.providers.standard.hooks.filesystem import FSHook as FSHook
42+
from airflow.providers.standard.hooks.package_index import PackageIndexHook as PackageIndexHook
43+
from airflow.providers.standard.hooks.subprocess import SubprocessHook as SubprocessHook
44+
from airflow.providers.standard.operators.bash import BashOperator as BashOperator
45+
from airflow.providers.standard.operators.branch import (
46+
BaseBranchOperator as BaseBranchOperator,
47+
BranchMixIn as BranchMixIn,
48+
)
49+
from airflow.providers.standard.operators.datetime import BranchDateTimeOperator as BranchDateTimeOperator
50+
from airflow.providers.standard.operators.empty import EmptyOperator as EmptyOperator
51+
from airflow.providers.standard.operators.latest_only import LatestOnlyOperator as LatestOnlyOperator
52+
from airflow.providers.standard.operators.python import (
53+
_SERIALIZERS as _SERIALIZERS,
54+
BranchExternalPythonOperator as BranchExternalPythonOperator,
55+
BranchPythonOperator as BranchPythonOperator,
56+
BranchPythonVirtualenvOperator as BranchPythonVirtualenvOperator,
57+
ExternalPythonOperator as ExternalPythonOperator,
58+
PythonOperator as PythonOperator,
59+
PythonVirtualenvOperator as PythonVirtualenvOperator,
60+
ShortCircuitOperator as ShortCircuitOperator,
61+
)
62+
from airflow.providers.standard.operators.smooth import SmoothOperator as SmoothOperator
63+
from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator as TriggerDagRunOperator
64+
from airflow.providers.standard.operators.weekday import BranchDayOfWeekOperator as BranchDayOfWeekOperator
65+
from airflow.providers.standard.sensors.bash import BashSensor as BashSensor
66+
from airflow.providers.standard.sensors.date_time import (
67+
DateTimeSensor as DateTimeSensor,
68+
DateTimeSensorAsync as DateTimeSensorAsync,
69+
)
70+
from airflow.providers.standard.sensors.external_task import (
71+
ExternalDagLink as ExternalDagLink,
72+
ExternalTaskMarker as ExternalTaskMarker,
73+
ExternalTaskSensor as ExternalTaskSensor,
74+
)
75+
from airflow.providers.standard.sensors.filesystem import FileSensor as FileSensor
76+
from airflow.providers.standard.sensors.python import PythonSensor as PythonSensor
77+
from airflow.providers.standard.sensors.time import (
78+
TimeSensor as TimeSensor,
79+
TimeSensorAsync as TimeSensorAsync,
80+
)
81+
from airflow.providers.standard.sensors.time_delta import (
82+
TimeDeltaSensor as TimeDeltaSensor,
83+
TimeDeltaSensorAsync as TimeDeltaSensorAsync,
84+
)
85+
from airflow.providers.standard.sensors.weekday import DayOfWeekSensor as DayOfWeekSensor
86+
from airflow.providers.standard.triggers.temporal import TimeDeltaTrigger as TimeDeltaTrigger
87+
from airflow.providers.standard.utils.python_virtualenv import (
88+
prepare_virtualenv as prepare_virtualenv,
89+
write_python_script as write_python_script,
90+
)
91+
from airflow.sdk import (
92+
DAG as DAG,
93+
Asset as Asset,
94+
AssetAlias as AssetAlias,
95+
AssetAll as AssetAll,
96+
AssetAny as AssetAny,
97+
BaseHook as BaseHook,
98+
BaseNotifier as BaseNotifier,
99+
BaseOperator as BaseOperator,
100+
BaseOperatorLink as BaseOperatorLink,
101+
BaseSensorOperator as BaseSensorOperator,
102+
Connection as Connection,
103+
Context as Context,
104+
DagRunState as DagRunState,
105+
EdgeModifier as EdgeModifier,
106+
Label as Label,
107+
Metadata as Metadata,
108+
ObjectStoragePath as ObjectStoragePath,
109+
Param as Param,
110+
PokeReturnValue as PokeReturnValue,
111+
TaskGroup as TaskGroup,
112+
TaskInstanceState as TaskInstanceState,
113+
TriggerRule as TriggerRule,
114+
Variable as Variable,
115+
WeightRule as WeightRule,
116+
XComArg as XComArg,
117+
chain as chain,
118+
chain_linear as chain_linear,
119+
cross_downstream as cross_downstream,
120+
dag as dag,
121+
get_current_context as get_current_context,
122+
get_parsing_context as get_parsing_context,
123+
setup as setup,
124+
task as task,
125+
task_group as task_group,
126+
teardown as teardown,
127+
)
128+
from airflow.sdk.bases.decorator import (
129+
DecoratedMappedOperator as DecoratedMappedOperator,
130+
DecoratedOperator as DecoratedOperator,
131+
TaskDecorator as TaskDecorator,
132+
)
133+
from airflow.sdk.bases.sensor import poke_mode_only as poke_mode_only
134+
from airflow.sdk.definitions.template import literal as literal
135+
from airflow.sdk.execution_time.xcom import XCom as XCom
136+
137+
__all__: list[str] = [
138+
"AirflowBadRequest",
139+
"AirflowConfigException",
140+
"AirflowException",
141+
"AirflowFailException",
142+
"AirflowNotFoundException",
143+
"AirflowSensorTimeout",
144+
"AirflowSkipException",
145+
"AirflowTaskTerminated",
146+
"AirflowTaskTimeout",
147+
"Asset",
148+
"AssetAlias",
149+
"AssetAll",
150+
"AssetAny",
151+
"BaseBranchOperator",
152+
"BaseHook",
153+
"BaseNotifier",
154+
"BaseOperator",
155+
"BaseOperatorLink",
156+
"BaseSensorOperator",
157+
"BashOperator",
158+
"BashSensor",
159+
"BranchDateTimeOperator",
160+
"BranchDayOfWeekOperator",
161+
"BranchExternalPythonOperator",
162+
"BranchMixIn",
163+
"BranchPythonOperator",
164+
"BranchPythonVirtualenvOperator",
165+
"Connection",
166+
"Context",
167+
"DAG",
168+
"DagRun",
169+
"DagRunState",
170+
"DateTimeSensor",
171+
"DateTimeSensorAsync",
172+
"DayOfWeekSensor",
173+
"DecoratedMappedOperator",
174+
"DecoratedOperator",
175+
"EdgeModifier",
176+
"EmptyOperator",
177+
"ExternalDagLink",
178+
"ExternalPythonOperator",
179+
"ExternalTaskMarker",
180+
"ExternalTaskSensor",
181+
"FSHook",
182+
"FileSensor",
183+
"Label",
184+
"LatestOnlyOperator",
185+
"MappedOperator",
186+
"Metadata",
187+
"ObjectStoragePath",
188+
"PackageIndexHook",
189+
"Param",
190+
"PokeReturnValue",
191+
"PythonOperator",
192+
"PythonSensor",
193+
"PythonVirtualenvOperator",
194+
"ShortCircuitOperator",
195+
"SmoothOperator",
196+
"SubprocessHook",
197+
"TaskDecorator",
198+
"TaskGroup",
199+
"TaskInstance",
200+
"TaskInstanceState",
201+
"TimeDeltaSensor",
202+
"TimeDeltaSensorAsync",
203+
"TimeDeltaTrigger",
204+
"TimeSensor",
205+
"TimeSensorAsync",
206+
"TriggerDagRunOperator",
207+
"TriggerRule",
208+
"Variable",
209+
"WeightRule",
210+
"XCOM_RETURN_KEY",
211+
"XCom",
212+
"XComArg",
213+
"_SERIALIZERS",
214+
"chain",
215+
"chain_linear",
216+
"cross_downstream",
217+
"dag",
218+
"get_current_context",
219+
"get_parsing_context",
220+
"io",
221+
"literal",
222+
"poke_mode_only",
223+
"prepare_virtualenv",
224+
"setup",
225+
"task",
226+
"task_group",
227+
"teardown",
228+
"timezone",
229+
"write_python_script",
230+
]

‎providers/common/compat/src/airflow/providers/common/compat/standard/operators.py‎

Lines changed: 8 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -17,35 +17,13 @@
1717

1818
from __future__ import annotations
1919

20-
from typing import TYPE_CHECKING
21-
22-
if TYPE_CHECKING:
23-
from airflow.providers.standard.operators.python import (
24-
_SERIALIZERS,
25-
PythonOperator,
26-
ShortCircuitOperator,
27-
get_current_context,
28-
)
29-
else:
30-
try:
31-
from airflow.providers.standard.operators.python import (
32-
_SERIALIZERS,
33-
PythonOperator,
34-
ShortCircuitOperator,
35-
get_current_context,
36-
)
37-
except ModuleNotFoundError:
38-
from airflow.operators.python import (
39-
_SERIALIZERS,
40-
PythonOperator,
41-
ShortCircuitOperator,
42-
)
43-
44-
try:
45-
from airflow.sdk import get_current_context
46-
except (ImportError, ModuleNotFoundError):
47-
from airflow.providers.standard.operators.python import get_current_context
48-
49-
from airflow.providers.common.compat.version_compat import BaseOperator
20+
# Re-export from lazy_compat for backward compatibility
21+
from airflow.providers.common.compat.lazy_compat import (
22+
_SERIALIZERS,
23+
BaseOperator,
24+
PythonOperator,
25+
ShortCircuitOperator,
26+
get_current_context,
27+
)
5028

5129
__all__ = ["BaseOperator", "PythonOperator", "_SERIALIZERS", "ShortCircuitOperator", "get_current_context"]

‎providers/common/compat/src/airflow/providers/common/compat/standard/triggers.py‎

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,7 @@
1717

1818
from __future__ import annotations
1919

20-
from typing import TYPE_CHECKING
21-
22-
if TYPE_CHECKING:
23-
from airflow.providers.standard.triggers.temporal import TimeDeltaTrigger
24-
else:
25-
try:
26-
from airflow.providers.standard.triggers.temporal import TimeDeltaTrigger
27-
except ModuleNotFoundError:
28-
from airflow.triggers.temporal import TimeDeltaTrigger
29-
20+
# Re-export from lazy_compat for backward compatibility
21+
from airflow.providers.common.compat.lazy_compat import TimeDeltaTrigger
3022

3123
__all__ = ["TimeDeltaTrigger"]

‎providers/common/compat/src/airflow/providers/common/compat/standard/utils.py‎

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,7 @@
1717

1818
from __future__ import annotations
1919

20-
from typing import TYPE_CHECKING
21-
22-
if TYPE_CHECKING:
23-
from airflow.providers.standard.utils.python_virtualenv import prepare_virtualenv, write_python_script
24-
else:
25-
try:
26-
from airflow.providers.standard.utils.python_virtualenv import prepare_virtualenv, write_python_script
27-
except ModuleNotFoundError:
28-
from airflow.utils.python_virtualenv import prepare_virtualenv, write_python_script
29-
20+
# Re-export from lazy_compat for backward compatibility
21+
from airflow.providers.common.compat.lazy_compat import prepare_virtualenv, write_python_script
3022

3123
__all__ = ["write_python_script", "prepare_virtualenv"]

‎providers/common/compat/tests/unit/common/compat/test_lazy_compat.py‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
#!/usr/bin/env python3
2+
# Licensed to the Apache Software Foundation (ASF) under one
3+
# or more contributor license agreements. See the NOTICE file
4+
# distributed with this work for additional information
5+
# regarding copyright ownership. The ASF licenses this file
6+
# to you under the Apache License, Version 2.0 (the
7+
# "License"); you may not use this file except in compliance
8+
# with the License. You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing,
13+
# software distributed under the License is distributed on an
14+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
# KIND, either express or implied. See the License for the
16+
# specific language governing permissions and limitations
17+
# under the License.
18+
19+
from __future__ import annotations
20+
21+
import pytest
22+
23+
24+
def test_all_compat_imports_work():
25+
"""
26+
Test that all items in _IMPORT_MAP can be successfully imported.
27+
28+
For each item, validates that at least one of the specified import paths works,
29+
ensuring the fallback mechanism is functional.
30+
"""
31+
from airflow.providers.common.compat import lazy_compat
32+
33+
failed_imports = []
34+
35+
for name in lazy_compat.__all__:
36+
try:
37+
obj = getattr(lazy_compat, name)
38+
assert obj is not None, f"{name} imported as None"
39+
except (ImportError, AttributeError) as e:
40+
failed_imports.append((name, str(e)))
41+
42+
if failed_imports:
43+
error_msg = "The following imports failed:\n"
44+
for name, error in failed_imports:
45+
error_msg += f" - {name}: {error}\n"
46+
pytest.fail(error_msg)
47+
48+
49+
def test_invalid_import_raises_attribute_error():
50+
"""Test that importing non-existent attribute raises AttributeError."""
51+
from airflow.providers.common.compat import lazy_compat
52+
53+
with pytest.raises(AttributeError, match="has no attribute 'NonExistentClass'"):
54+
_ = lazy_compat.NonExistentClass

0 commit comments

Comments
 (0)