diff --git a/providers/grpc/src/airflow/providers/grpc/hooks/grpc.py b/providers/grpc/src/airflow/providers/grpc/hooks/grpc.py index 7395fcd2a42d2..2cb1fc93b625b 100644 --- a/providers/grpc/src/airflow/providers/grpc/hooks/grpc.py +++ b/providers/grpc/src/airflow/providers/grpc/hooks/grpc.py @@ -30,11 +30,7 @@ ) from airflow.exceptions import AirflowConfigException - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.grpc.version_compat import BaseHook class GrpcHook(BaseHook): diff --git a/providers/grpc/src/airflow/providers/grpc/version_compat.py b/providers/grpc/src/airflow/providers/grpc/version_compat.py index ab72155414178..c98ecb412aa57 100644 --- a/providers/grpc/src/airflow/providers/grpc/version_compat.py +++ b/providers/grpc/src/airflow/providers/grpc/version_compat.py @@ -28,10 +28,16 @@ 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 else: from airflow.models import BaseOperator # type: ignore[no-redef] -__all__ = ["AIRFLOW_V_3_0_PLUS", "BaseOperator"] +__all__ = ["AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", "BaseHook", "BaseOperator"]