From 7b7402a7e5aeb192387d599a5e1dccbebeb23b4b Mon Sep 17 00:00:00 2001
From: Dov Benyomin Sohacheski <b@kloud.email>
Date: Thu, 3 Jul 2025 17:22:10 +0300
Subject: [PATCH] Replace BaseHook to Task SDK for apache/beam

---
 .../beam/src/airflow/providers/apache/beam/hooks/beam.py  | 8 +-------
 .../src/airflow/providers/apache/beam/version_compat.py   | 4 +++-
 2 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/providers/apache/beam/src/airflow/providers/apache/beam/hooks/beam.py b/providers/apache/beam/src/airflow/providers/apache/beam/hooks/beam.py
index 635e9b84d82d7..e40e7f086e780 100644
--- a/providers/apache/beam/src/airflow/providers/apache/beam/hooks/beam.py
+++ b/providers/apache/beam/src/airflow/providers/apache/beam/hooks/beam.py
@@ -15,8 +15,6 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""This module contains a Apache Beam Hook."""
-
 from __future__ import annotations
 
 import asyncio
@@ -37,13 +35,9 @@
 from packaging.version import Version
 
 from airflow.exceptions import AirflowConfigException, AirflowException
+from airflow.providers.apache.beam.version_compat import BaseHook
 from airflow.providers.common.compat.standard.utils import prepare_virtualenv
 
-try:
-    from airflow.sdk import BaseHook
-except ImportError:
-    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
-
 if TYPE_CHECKING:
     import logging
 
diff --git a/providers/apache/beam/src/airflow/providers/apache/beam/version_compat.py b/providers/apache/beam/src/airflow/providers/apache/beam/version_compat.py
index 6756c5b297327..9f69b38930166 100644
--- a/providers/apache/beam/src/airflow/providers/apache/beam/version_compat.py
+++ b/providers/apache/beam/src/airflow/providers/apache/beam/version_compat.py
@@ -32,11 +32,13 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]:
 AIRFLOW_V_3_1_PLUS = get_base_airflow_version_tuple() >= (3, 1, 0)
 
 if AIRFLOW_V_3_1_PLUS:
-    from airflow.sdk import BaseOperator
+    from airflow.sdk import BaseHook, BaseOperator
 else:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
     from airflow.models import BaseOperator
 
 __all__ = [
     "AIRFLOW_V_3_1_PLUS",
+    "BaseHook",
     "BaseOperator",
 ]