Skip to content

Commit 8f8ddf8

Browse files
authored
Move all BaseHook usages to version_compat in MongoDB (#52857)
Part of #52676
1 parent 6654c15 commit 8f8ddf8

File tree

2 files changed

+8
-6
lines changed

2 files changed

+8
-6
lines changed

providers/mongo/src/airflow/providers/mongo/hooks/mongo.py

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,11 +28,7 @@
2828
from pymongo.errors import CollectionInvalid
2929

3030
from airflow.exceptions import AirflowConfigException
31-
32-
try:
33-
from airflow.sdk import BaseHook
34-
except ImportError:
35-
from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef]
31+
from airflow.providers.mongo.version_compat import BaseHook
3632

3733
if TYPE_CHECKING:
3834
from types import TracebackType

providers/mongo/src/airflow/providers/mongo/version_compat.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,10 +33,16 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]:
3333

3434

3535
AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0)
36+
AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0)
37+
38+
if AIRFLOW_V_3_1_PLUS:
39+
from airflow.sdk import BaseHook
40+
else:
41+
from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef]
3642

3743
if AIRFLOW_V_3_0_PLUS:
3844
from airflow.sdk import BaseSensorOperator
3945
else:
4046
from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef]
4147

42-
__all__ = ["AIRFLOW_V_3_0_PLUS", "BaseSensorOperator"]
48+
__all__ = ["AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", "BaseHook", "BaseSensorOperator"]

0 commit comments

Comments
 (0)