From f2404c93628bbb3d10e7537551e99620b2a08429 Mon Sep 17 00:00:00 2001
From: Ankit Chaurasia <8670962+sunank200@users.noreply.github.com>
Date: Thu, 3 Jul 2025 22:04:29 +0545
Subject: [PATCH] Move all BaseHook usages to version_compat in Docker

---
 .../src/airflow/providers/docker/hooks/docker.py     |  6 +-----
 .../src/airflow/providers/docker/version_compat.py   | 12 +++++++++++-
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/providers/docker/src/airflow/providers/docker/hooks/docker.py b/providers/docker/src/airflow/providers/docker/hooks/docker.py
index c07ba908392e6..452a45b83fb20 100644
--- a/providers/docker/src/airflow/providers/docker/hooks/docker.py
+++ b/providers/docker/src/airflow/providers/docker/hooks/docker.py
@@ -27,11 +27,7 @@
 from docker.errors import APIError, DockerException
 
 from airflow.exceptions import AirflowException, AirflowNotFoundException
-
-try:
-    from airflow.sdk import BaseHook
-except ImportError:
-    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
+from airflow.providers.docker.version_compat import BaseHook
 
 if TYPE_CHECKING:
     from airflow.models import Connection
diff --git a/providers/docker/src/airflow/providers/docker/version_compat.py b/providers/docker/src/airflow/providers/docker/version_compat.py
index f2ac9959a5222..615e9c2f23db7 100644
--- a/providers/docker/src/airflow/providers/docker/version_compat.py
+++ b/providers/docker/src/airflow/providers/docker/version_compat.py
@@ -33,11 +33,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_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
 
-__all__ = ["AIRFLOW_V_3_0_PLUS", "BaseOperator"]
+__all__ = [
+    "AIRFLOW_V_3_0_PLUS",
+    "AIRFLOW_V_3_1_PLUS",
+    "BaseHook",
+    "BaseOperator",
+]