diff --git a/providers/docker/pyproject.toml b/providers/docker/pyproject.toml index 3454bb9da070a..773ed98fa38d7 100644 --- a/providers/docker/pyproject.toml +++ b/providers/docker/pyproject.toml @@ -58,6 +58,7 @@ requires-python = ">=3.10" # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build`` dependencies = [ "apache-airflow>=2.10.0", + "apache-airflow-providers-common-compat>=1.8.0", "docker>=7.1.0", "python-dotenv>=0.21.0", ] diff --git a/providers/docker/src/airflow/providers/docker/decorators/docker.py b/providers/docker/src/airflow/providers/docker/decorators/docker.py index 807902d5e8170..c64c9560b592b 100644 --- a/providers/docker/src/airflow/providers/docker/decorators/docker.py +++ b/providers/docker/src/airflow/providers/docker/decorators/docker.py @@ -35,14 +35,9 @@ if TYPE_CHECKING: from typing import Literal + from airflow.providers.common.compat.sdk import Context from airflow.sdk.bases.decorator import TaskDecorator - if AIRFLOW_V_3_0_PLUS: - from airflow.sdk.definitions.context import Context - else: - # TODO: Remove once provider drops support for Airflow 2 - from airflow.utils.context import Context - Serializer = Literal["pickle", "dill", "cloudpickle"] try: diff --git a/providers/docker/src/airflow/providers/docker/hooks/docker.py b/providers/docker/src/airflow/providers/docker/hooks/docker.py index 452a45b83fb20..5f1a02f5830c3 100644 --- a/providers/docker/src/airflow/providers/docker/hooks/docker.py +++ b/providers/docker/src/airflow/providers/docker/hooks/docker.py @@ -27,10 +27,10 @@ from docker.errors import APIError, DockerException from airflow.exceptions import AirflowException, AirflowNotFoundException -from airflow.providers.docker.version_compat import BaseHook +from airflow.providers.common.compat.sdk import BaseHook if TYPE_CHECKING: - from airflow.models import Connection + from airflow.providers.common.compat.sdk import Connection class DockerHook(BaseHook): diff --git a/providers/docker/src/airflow/providers/docker/operators/docker.py b/providers/docker/src/airflow/providers/docker/operators/docker.py index dc8065d79cc6d..045c82d807a8f 100644 --- a/providers/docker/src/airflow/providers/docker/operators/docker.py +++ b/providers/docker/src/airflow/providers/docker/operators/docker.py @@ -45,7 +45,7 @@ from docker import APIClient from docker.types import DeviceRequest - from airflow.sdk.definitions.context import Context + from airflow.providers.common.compat.sdk import Context from airflow.sdk.types import Logger diff --git a/providers/docker/src/airflow/providers/docker/operators/docker_swarm.py b/providers/docker/src/airflow/providers/docker/operators/docker_swarm.py index b5a3119b37d2f..dea49f36f2e39 100644 --- a/providers/docker/src/airflow/providers/docker/operators/docker_swarm.py +++ b/providers/docker/src/airflow/providers/docker/operators/docker_swarm.py @@ -32,11 +32,7 @@ from airflow.utils.strings import get_random_string if TYPE_CHECKING: - try: - from airflow.sdk.definitions.context import Context - except ImportError: - # TODO: Remove once provider drops support for Airflow 2 - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class DockerSwarmOperator(DockerOperator): diff --git a/providers/docker/src/airflow/providers/docker/version_compat.py b/providers/docker/src/airflow/providers/docker/version_compat.py index 9ce966132fd6c..29c46806a0a6e 100644 --- a/providers/docker/src/airflow/providers/docker/version_compat.py +++ b/providers/docker/src/airflow/providers/docker/version_compat.py @@ -40,16 +40,13 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: # This is needed for DecoratedOperator compatibility if AIRFLOW_V_3_1_PLUS: from airflow.sdk import ( - BaseHook, BaseOperator, ) else: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] from airflow.models import BaseOperator __all__ = [ "AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", - "BaseHook", "BaseOperator", ]