From 008cf226cf9829a2c5a3b81fa1884327cdbcd915 Mon Sep 17 00:00:00 2001
From: Ankit Chaurasia <8670962+sunank200@users.noreply.github.com>
Date: Fri, 4 Jul 2025 15:39:28 +0545
Subject: [PATCH] Move all BaseHook usages to version_compat in Microsoft WinRM

Part of https://github.com/apache/airflow/issues/52676
---
 .../src/airflow/providers/microsoft/winrm/hooks/winrm.py  | 6 +-----
 .../airflow/providers/microsoft/winrm/version_compat.py   | 8 ++++++++
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/providers/microsoft/winrm/src/airflow/providers/microsoft/winrm/hooks/winrm.py b/providers/microsoft/winrm/src/airflow/providers/microsoft/winrm/hooks/winrm.py
index 75efdc1588c36..4daf5ca09222e 100644
--- a/providers/microsoft/winrm/src/airflow/providers/microsoft/winrm/hooks/winrm.py
+++ b/providers/microsoft/winrm/src/airflow/providers/microsoft/winrm/hooks/winrm.py
@@ -26,11 +26,7 @@
 from winrm.protocol import Protocol
 
 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.microsoft.winrm.version_compat import BaseHook
 from airflow.utils.platform import getuser
 
 # TODO: FIXME please - I have too complex implementation
diff --git a/providers/microsoft/winrm/src/airflow/providers/microsoft/winrm/version_compat.py b/providers/microsoft/winrm/src/airflow/providers/microsoft/winrm/version_compat.py
index 89ade07a76753..497d0a9c8b7bb 100644
--- a/providers/microsoft/winrm/src/airflow/providers/microsoft/winrm/version_compat.py
+++ b/providers/microsoft/winrm/src/airflow/providers/microsoft/winrm/version_compat.py
@@ -28,13 +28,21 @@ 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_0_PLUS:
     from airflow.sdk import BaseOperator
 else:
     from airflow.models import BaseOperator  # type: ignore[no-redef]
 
+if AIRFLOW_V_3_1_PLUS:
+    from airflow.sdk import BaseHook
+else:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
+
 __all__ = [
     "AIRFLOW_V_3_0_PLUS",
+    "AIRFLOW_V_3_1_PLUS",
     "BaseOperator",
+    "BaseHook",
 ]