From 4263175351983b8b27ed2cd3903442c2b4ae559d Mon Sep 17 00:00:00 2001
From: abhishekbhakat <abhishek.bhakat@hotmail.com>
Date: Sun, 4 May 2025 09:12:15 +0000
Subject: [PATCH] feat: add version check for Airflow tools and remove client
 initialization function

---
 src/airflow_mcp_server/tools/tool_manager.py | 30 +++++++++-----------
 1 file changed, 14 insertions(+), 16 deletions(-)

diff --git a/src/airflow_mcp_server/tools/tool_manager.py b/src/airflow_mcp_server/tools/tool_manager.py
index d837a61..5079cef 100644
--- a/src/airflow_mcp_server/tools/tool_manager.py
+++ b/src/airflow_mcp_server/tools/tool_manager.py
@@ -1,6 +1,7 @@
 import logging
 
 from mcp.types import Tool
+from packaging.version import parse as parse_version
 
 from airflow_mcp_server.client.airflow_client import AirflowClient
 from airflow_mcp_server.config import AirflowConfig
@@ -12,22 +13,6 @@ logger = logging.getLogger(__name__)
 _tools_cache: dict[str, AirflowTool] = {}
 
 
-def _initialize_client(config: AirflowConfig) -> AirflowClient:
-    """Initialize Airflow client with configuration.
-
-    Args:
-        config: Configuration object with auth and URL settings
-
-    Returns:
-        AirflowClient instance
-
-    Raises:
-        ValueError: If default spec is not found
-    """
-    # Only use base_url and auth_token
-    return AirflowClient(base_url=config.base_url, auth_token=config.auth_token)
-
-
 async def _initialize_tools(config: AirflowConfig) -> None:
     """Initialize tools cache with Airflow operations (async)."""
     global _tools_cache
@@ -58,9 +43,22 @@ async def get_airflow_tools(config: AirflowConfig, mode: str = "unsafe") -> list
     Raises:
         ValueError: If initialization fails
     """
+
+    # Version check before returning tools
     if not _tools_cache:
         await _initialize_tools(config)
 
+    # Only check version if get_version tool is present
+    if "get_version" in _tools_cache:
+        version_tool = _tools_cache["get_version"]
+        async with version_tool.client:
+            version_result = await version_tool.run()
+        airflow_version = version_result.get("version")
+        if airflow_version is None:
+            raise RuntimeError("Could not determine Airflow version from get_version tool.")
+        if parse_version(airflow_version) <= parse_version("3.1.0"):
+            raise RuntimeError(f"Airflow version {airflow_version} is not supported. Requires > 3.1.0.")
+
     tools = []
     for operation_id, tool in _tools_cache.items():
         try: