From 3391fc65e4e2b734ff970030dec276649cbf3622 Mon Sep 17 00:00:00 2001 From: Ankit Chaurasia <8670962+sunank200@users.noreply.github.com> Date: Thu, 3 Jul 2025 22:42:43 +0545 Subject: [PATCH] Move all BaseHook usages to version_compat in Github Part of https://github.com/apache/airflow/issues/52676 --- .../github/src/airflow/providers/github/hooks/github.py | 6 +----- .../github/src/airflow/providers/github/version_compat.py | 8 ++++++++ 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/providers/github/src/airflow/providers/github/hooks/github.py b/providers/github/src/airflow/providers/github/hooks/github.py index d833be9d4b38b..b964f597140aa 100644 --- a/providers/github/src/airflow/providers/github/hooks/github.py +++ b/providers/github/src/airflow/providers/github/hooks/github.py @@ -24,11 +24,7 @@ from github import Github as GithubClient 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.github.version_compat import BaseHook class GithubHook(BaseHook): diff --git a/providers/github/src/airflow/providers/github/version_compat.py b/providers/github/src/airflow/providers/github/version_compat.py index 682d19a5f49be..7a03eaab7d19d 100644 --- a/providers/github/src/airflow/providers/github/version_compat.py +++ b/providers/github/src/airflow/providers/github/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 @@ -45,6 +51,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",