feat: add version check for Airflow tools and remove client initialization function

This commit is contained in:
2025-05-04 09:12:15 +00:00
parent b5cf563b8f
commit 4263175351

View File

@@ -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: