diff --git a/providers/mongo/pyproject.toml b/providers/mongo/pyproject.toml index 8c7f010036d4e..4508bdda07f1b 100644 --- a/providers/mongo/pyproject.toml +++ b/providers/mongo/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", "dnspython>=1.13.0", "pymongo>=4.13.2", ] @@ -67,6 +68,7 @@ dev = [ "apache-airflow", "apache-airflow-task-sdk", "apache-airflow-devel-common", + "apache-airflow-providers-common-compat", # Additional devel dependencies (do not remove this line and add extra development dependencies) "testcontainers>=4.12.0" ] diff --git a/providers/mongo/src/airflow/providers/mongo/hooks/mongo.py b/providers/mongo/src/airflow/providers/mongo/hooks/mongo.py index 64e747fe2fd22..d1f392ffdea52 100644 --- a/providers/mongo/src/airflow/providers/mongo/hooks/mongo.py +++ b/providers/mongo/src/airflow/providers/mongo/hooks/mongo.py @@ -28,7 +28,7 @@ from pymongo.errors import CollectionInvalid from airflow.exceptions import AirflowConfigException -from airflow.providers.mongo.version_compat import BaseHook +from airflow.providers.common.compat.sdk import BaseHook if TYPE_CHECKING: from types import TracebackType diff --git a/providers/mongo/src/airflow/providers/mongo/sensors/mongo.py b/providers/mongo/src/airflow/providers/mongo/sensors/mongo.py index 56d6d60bc7b7e..b8b152b49b4bf 100644 --- a/providers/mongo/src/airflow/providers/mongo/sensors/mongo.py +++ b/providers/mongo/src/airflow/providers/mongo/sensors/mongo.py @@ -20,8 +20,8 @@ from collections.abc import Sequence from typing import TYPE_CHECKING +from airflow.providers.common.compat.sdk import BaseSensorOperator from airflow.providers.mongo.hooks.mongo import MongoHook -from airflow.providers.mongo.version_compat import BaseSensorOperator if TYPE_CHECKING: try: diff --git a/providers/mongo/src/airflow/providers/mongo/version_compat.py b/providers/mongo/src/airflow/providers/mongo/version_compat.py index 2002927fcf990..a4fae50317e81 100644 --- a/providers/mongo/src/airflow/providers/mongo/version_compat.py +++ b/providers/mongo/src/airflow/providers/mongo/version_compat.py @@ -35,14 +35,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0) -if AIRFLOW_V_3_1_PLUS: - from airflow.sdk import BaseHook -else: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseSensorOperator -else: - from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef] - -__all__ = ["AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", "BaseHook", "BaseSensorOperator"] +__all__ = [ + "AIRFLOW_V_3_0_PLUS", + "AIRFLOW_V_3_1_PLUS", +]