diff --git a/providers/sftp/src/airflow/providers/sftp/hooks/sftp.py b/providers/sftp/src/airflow/providers/sftp/hooks/sftp.py index 540150f453046..7064bc923d427 100644 --- a/providers/sftp/src/airflow/providers/sftp/hooks/sftp.py +++ b/providers/sftp/src/airflow/providers/sftp/hooks/sftp.py @@ -35,13 +35,9 @@ from asgiref.sync import sync_to_async from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning +from airflow.providers.sftp.version_compat import BaseHook from airflow.providers.ssh.hooks.ssh import SSHHook -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] - if TYPE_CHECKING: from paramiko import SSHClient from paramiko.sftp_attr import SFTPAttributes diff --git a/providers/sftp/src/airflow/providers/sftp/version_compat.py b/providers/sftp/src/airflow/providers/sftp/version_compat.py index cb0e95874b6be..06f57818ab077 100644 --- a/providers/sftp/src/airflow/providers/sftp/version_compat.py +++ b/providers/sftp/src/airflow/providers/sftp/version_compat.py @@ -33,6 +33,12 @@ 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 BaseOperator, BaseSensorOperator, PokeReturnValue @@ -40,4 +46,4 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: from airflow.models import BaseOperator from airflow.sensors.base import BaseSensorOperator, PokeReturnValue # type: ignore[no-redef] -__all__ = ["AIRFLOW_V_3_0_PLUS", "BaseOperator", "BaseSensorOperator", "PokeReturnValue"] +__all__ = ["AIRFLOW_V_3_0_PLUS", "BaseHook", "BaseOperator", "BaseSensorOperator", "PokeReturnValue"]