diff --git a/providers/airbyte/src/airflow/providers/airbyte/hooks/airbyte.py b/providers/airbyte/src/airflow/providers/airbyte/hooks/airbyte.py index dc7c39683176d..5c30b1ddef32f 100644 --- a/providers/airbyte/src/airflow/providers/airbyte/hooks/airbyte.py +++ b/providers/airbyte/src/airflow/providers/airbyte/hooks/airbyte.py @@ -26,11 +26,7 @@ from requests import Session from airflow.exceptions import AirflowException - -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.airbyte.version_compat import BaseHook T = TypeVar("T", bound=Any) diff --git a/providers/airbyte/src/airflow/providers/airbyte/version_compat.py b/providers/airbyte/src/airflow/providers/airbyte/version_compat.py index 5faab587ffcef..da458b9effc70 100644 --- a/providers/airbyte/src/airflow/providers/airbyte/version_compat.py +++ b/providers/airbyte/src/airflow/providers/airbyte/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 @@ -44,6 +50,8 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: __all__ = [ "AIRFLOW_V_3_0_PLUS", + "AIRFLOW_V_3_1_PLUS", + "BaseHook", "BaseOperator", "BaseSensorOperator", "Context",