Skip to content

Commit 6865035

Browse files
authored
Move all BaseHook usages in providers to version_compat in ssh (#52780)
1 parent b51476b commit 6865035

File tree

2 files changed

+8
-6
lines changed

2 files changed

+8
-6
lines changed

providers/ssh/src/airflow/providers/ssh/hooks/ssh.py

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -33,11 +33,7 @@
3333
from tenacity import Retrying, stop_after_attempt, wait_fixed, wait_random
3434

3535
from airflow.exceptions import AirflowException
36-
37-
try:
38-
from airflow.sdk import BaseHook
39-
except ImportError:
40-
from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef]
36+
from airflow.providers.ssh.version_compat import BaseHook
4137
from airflow.utils.platform import getuser
4238
from airflow.utils.types import NOTSET, ArgNotSet
4339

providers/ssh/src/airflow/providers/ssh/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 BaseOperator
3945
else:
4046
from airflow.models import BaseOperator # type: ignore[no-redef]
4147

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

0 commit comments

Comments
 (0)