Skip to content

Commit a35602c

Browse files
authored
Simplify version-specific imports in the Standard provider (apache#56867)
Consolidate the Standard provider's compatibility logic by replacing imports from version_compat.py with direct imports from lazy_compat, providing a single source of truth for all Airflow 2.x ↔ 3.x compatibility. Follow-up of apache#56790 & apache#56793
1 parent 652dd85 commit a35602c

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

47 files changed

+263
-757
lines changed

.pre-commit-config.yaml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -466,9 +466,9 @@ repos:
466466
--pattern '^openlineage\.client\.(facet|run)'
467467
--message "You should import from `airflow.providers.common.compat.openlineage.facet` instead."
468468
- id: check-common-compat-sdk-imports-in-sync
469-
name: Check common.compat sdk.pyi is in sync
469+
name: Check common.compat sdk TYPE_CHECKING matches runtime maps
470470
language: python
471-
files: ^providers/common/compat/src/airflow/providers/common/compat/sdk\.(py|pyi)$
471+
files: ^providers/common/compat/src/airflow/providers/common/compat/sdk\.py$
472472
pass_filenames: false
473473
entry: ./scripts/ci/prek/check_common_compat_lazy_imports.py
474474
- id: check-airflow-providers-bug-report-template

providers/common/compat/src/airflow/providers/common/compat/sdk.py

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,64 @@
2323

2424
from __future__ import annotations
2525

26+
from typing import TYPE_CHECKING
27+
28+
if TYPE_CHECKING:
29+
import airflow.sdk.io as io # noqa: F401
30+
import airflow.sdk.timezone as timezone # noqa: F401
31+
from airflow.models.xcom import XCOM_RETURN_KEY as XCOM_RETURN_KEY
32+
from airflow.sdk import (
33+
DAG as DAG,
34+
Asset as Asset,
35+
AssetAlias as AssetAlias,
36+
AssetAll as AssetAll,
37+
AssetAny as AssetAny,
38+
BaseHook as BaseHook,
39+
BaseNotifier as BaseNotifier,
40+
BaseOperator as BaseOperator,
41+
BaseOperatorLink as BaseOperatorLink,
42+
BaseSensorOperator as BaseSensorOperator,
43+
Connection as Connection,
44+
Context as Context,
45+
DagRunState as DagRunState,
46+
EdgeModifier as EdgeModifier,
47+
Label as Label,
48+
Metadata as Metadata,
49+
ObjectStoragePath as ObjectStoragePath,
50+
Param as Param,
51+
PokeReturnValue as PokeReturnValue,
52+
TaskGroup as TaskGroup,
53+
TaskInstanceState as TaskInstanceState,
54+
TriggerRule as TriggerRule,
55+
Variable as Variable,
56+
WeightRule as WeightRule,
57+
XComArg as XComArg,
58+
chain as chain,
59+
chain_linear as chain_linear,
60+
cross_downstream as cross_downstream,
61+
dag as dag,
62+
get_current_context as get_current_context,
63+
get_parsing_context as get_parsing_context,
64+
setup as setup,
65+
task as task,
66+
task_group as task_group,
67+
teardown as teardown,
68+
)
69+
from airflow.sdk.bases.decorator import (
70+
DecoratedMappedOperator as DecoratedMappedOperator,
71+
DecoratedOperator as DecoratedOperator,
72+
TaskDecorator as TaskDecorator,
73+
get_unique_task_id as get_unique_task_id,
74+
task_decorator_factory as task_decorator_factory,
75+
)
76+
from airflow.sdk.bases.sensor import poke_mode_only as poke_mode_only
77+
from airflow.sdk.definitions.context import context_merge as context_merge
78+
from airflow.sdk.definitions.mappedoperator import MappedOperator as MappedOperator
79+
from airflow.sdk.definitions.template import literal as literal
80+
from airflow.sdk.execution_time.context import context_to_airflow_vars as context_to_airflow_vars
81+
from airflow.sdk.execution_time.timeout import timeout as timeout
82+
from airflow.sdk.execution_time.xcom import XCom as XCom
83+
2684
from airflow.providers.common.compat._compat_utils import create_module_getattr
2785

2886
# Rename map for classes that changed names between Airflow 2.x and 3.x
@@ -63,6 +121,8 @@
63121
"setup": ("airflow.sdk", "airflow.decorators"),
64122
"teardown": ("airflow.sdk", "airflow.decorators"),
65123
"TaskDecorator": ("airflow.sdk.bases.decorator", "airflow.decorators"),
124+
"task_decorator_factory": ("airflow.sdk.bases.decorator", "airflow.decorators.base"),
125+
"get_unique_task_id": ("airflow.sdk.bases.decorator", "airflow.decorators.base"),
66126
# ============================================================================
67127
# Models
68128
# ============================================================================
@@ -74,6 +134,7 @@
74134
"XComArg": ("airflow.sdk", "airflow.models.xcom_arg"),
75135
"DecoratedOperator": ("airflow.sdk.bases.decorator", "airflow.decorators.base"),
76136
"DecoratedMappedOperator": ("airflow.sdk.bases.decorator", "airflow.decorators.base"),
137+
"MappedOperator": ("airflow.sdk.definitions.mappedoperator", "airflow.models.mappedoperator"),
77138
# ============================================================================
78139
# Assets (Dataset → Asset rename in Airflow 3.0)
79140
# ============================================================================
@@ -119,6 +180,8 @@
119180
# Context & Utilities
120181
# ============================================================================
121182
"Context": ("airflow.sdk", "airflow.utils.context"),
183+
"context_merge": ("airflow.sdk.definitions.context", "airflow.utils.context"),
184+
"context_to_airflow_vars": ("airflow.sdk.execution_time.context", "airflow.utils.operator_helpers"),
122185
"get_current_context": ("airflow.sdk", "airflow.operators.python"),
123186
"get_parsing_context": ("airflow.sdk", "airflow.utils.dag_parsing_context"),
124187
# ============================================================================

providers/common/compat/src/airflow/providers/common/compat/sdk.pyi

Lines changed: 0 additions & 120 deletions
This file was deleted.

providers/standard/pyproject.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,13 +58,15 @@ requires-python = ">=3.10"
5858
# After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build``
5959
dependencies = [
6060
"apache-airflow>=2.10.0",
61+
"apache-airflow-providers-common-compat>=1.7.4", # + TODO: bump to next version
6162
]
6263

6364
[dependency-groups]
6465
dev = [
6566
"apache-airflow",
6667
"apache-airflow-task-sdk",
6768
"apache-airflow-devel-common",
69+
"apache-airflow-providers-common-compat",
6870
# Additional devel dependencies (do not remove this line and add extra development dependencies)
6971
"apache-airflow-providers-mysql",
7072
]

providers/standard/src/airflow/providers/standard/decorators/bash.py

Lines changed: 7 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -21,24 +21,18 @@
2121
from collections.abc import Callable, Collection, Mapping, Sequence
2222
from typing import TYPE_CHECKING, Any, ClassVar
2323

24-
from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
25-
26-
if AIRFLOW_V_3_0_PLUS:
27-
from airflow.sdk.bases.decorator import DecoratedOperator, TaskDecorator, task_decorator_factory
28-
else:
29-
from airflow.decorators.base import ( # type: ignore[no-redef]
30-
DecoratedOperator,
31-
TaskDecorator,
32-
task_decorator_factory,
33-
)
34-
24+
from airflow.providers.common.compat.sdk import (
25+
DecoratedOperator,
26+
TaskDecorator,
27+
context_merge,
28+
task_decorator_factory,
29+
)
3530
from airflow.providers.standard.operators.bash import BashOperator
36-
from airflow.providers.standard.version_compat import context_merge
3731
from airflow.sdk.definitions._internal.types import SET_DURING_EXECUTION
3832
from airflow.utils.operator_helpers import determine_kwargs
3933

4034
if TYPE_CHECKING:
41-
from airflow.sdk.definitions.context import Context
35+
from airflow.providers.common.compat.sdk import Context
4236

4337

4438
class _BashDecoratedOperator(DecoratedOperator, BashOperator):

providers/standard/src/airflow/providers/standard/decorators/branch_external_python.py

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -19,18 +19,12 @@
1919
from collections.abc import Callable
2020
from typing import TYPE_CHECKING
2121

22-
from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
23-
24-
if AIRFLOW_V_3_0_PLUS:
25-
from airflow.sdk.bases.decorator import task_decorator_factory
26-
else:
27-
from airflow.decorators.base import task_decorator_factory # type: ignore[no-redef]
28-
22+
from airflow.providers.common.compat.sdk import task_decorator_factory
2923
from airflow.providers.standard.decorators.python import _PythonDecoratedOperator
3024
from airflow.providers.standard.operators.python import BranchExternalPythonOperator
3125

3226
if TYPE_CHECKING:
33-
from airflow.sdk.bases.decorator import TaskDecorator
27+
from airflow.providers.common.compat.sdk import TaskDecorator
3428

3529

3630
class _BranchExternalPythonDecoratedOperator(_PythonDecoratedOperator, BranchExternalPythonOperator):

providers/standard/src/airflow/providers/standard/decorators/branch_python.py

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,17 +19,12 @@
1919
from collections.abc import Callable
2020
from typing import TYPE_CHECKING
2121

22-
from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
23-
24-
if AIRFLOW_V_3_0_PLUS:
25-
from airflow.sdk.bases.decorator import task_decorator_factory
26-
else:
27-
from airflow.decorators.base import task_decorator_factory # type: ignore[no-redef]
22+
from airflow.providers.common.compat.sdk import task_decorator_factory
2823
from airflow.providers.standard.decorators.python import _PythonDecoratedOperator
2924
from airflow.providers.standard.operators.python import BranchPythonOperator
3025

3126
if TYPE_CHECKING:
32-
from airflow.sdk.bases.decorator import TaskDecorator
27+
from airflow.providers.common.compat.sdk import TaskDecorator
3328

3429

3530
class _BranchPythonDecoratedOperator(_PythonDecoratedOperator, BranchPythonOperator):

providers/standard/src/airflow/providers/standard/decorators/branch_virtualenv.py

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,17 +19,12 @@
1919
from collections.abc import Callable
2020
from typing import TYPE_CHECKING
2121

22-
from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
23-
24-
if AIRFLOW_V_3_0_PLUS:
25-
from airflow.sdk.bases.decorator import task_decorator_factory
26-
else:
27-
from airflow.decorators.base import task_decorator_factory # type: ignore[no-redef]
22+
from airflow.providers.common.compat.sdk import task_decorator_factory
2823
from airflow.providers.standard.decorators.python import _PythonDecoratedOperator
2924
from airflow.providers.standard.operators.python import BranchPythonVirtualenvOperator
3025

3126
if TYPE_CHECKING:
32-
from airflow.sdk.bases.decorator import TaskDecorator
27+
from airflow.providers.common.compat.sdk import TaskDecorator
3328

3429

3530
class _BranchPythonVirtualenvDecoratedOperator(_PythonDecoratedOperator, BranchPythonVirtualenvOperator):

providers/standard/src/airflow/providers/standard/decorators/external_python.py

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,17 +19,12 @@
1919
from collections.abc import Callable
2020
from typing import TYPE_CHECKING
2121

22-
from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
23-
24-
if AIRFLOW_V_3_0_PLUS:
25-
from airflow.sdk.bases.decorator import task_decorator_factory
26-
else:
27-
from airflow.decorators.base import task_decorator_factory # type: ignore[no-redef]
22+
from airflow.providers.common.compat.sdk import task_decorator_factory
2823
from airflow.providers.standard.decorators.python import _PythonDecoratedOperator
2924
from airflow.providers.standard.operators.python import ExternalPythonOperator
3025

3126
if TYPE_CHECKING:
32-
from airflow.sdk.bases.decorator import TaskDecorator
27+
from airflow.providers.common.compat.sdk import TaskDecorator
3328

3429

3530
class _PythonExternalDecoratedOperator(_PythonDecoratedOperator, ExternalPythonOperator):

providers/standard/src/airflow/providers/standard/decorators/python.py

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,16 +19,11 @@
1919
from collections.abc import Callable, Sequence
2020
from typing import TYPE_CHECKING
2121

22+
from airflow.providers.common.compat.sdk import DecoratedOperator, task_decorator_factory
2223
from airflow.providers.standard.operators.python import PythonOperator
23-
from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
24-
25-
if AIRFLOW_V_3_0_PLUS:
26-
from airflow.sdk.bases.decorator import DecoratedOperator, task_decorator_factory
27-
else:
28-
from airflow.decorators.base import DecoratedOperator, task_decorator_factory # type: ignore[no-redef]
2924

3025
if TYPE_CHECKING:
31-
from airflow.sdk.bases.decorator import TaskDecorator
26+
from airflow.providers.common.compat.sdk import TaskDecorator
3227

3328

3429
class _PythonDecoratedOperator(DecoratedOperator, PythonOperator):

0 commit comments

Comments
 (0)