Clean up for only MCP Server

src/airflow_mcp_server/tools/tool_manager.py

@@ -0,0 +1,126 @@
+import logging
+import os
+from importlib import resources
+
+from mcp.types import Tool
+
+from airflow_mcp_server.client.airflow_client import AirflowClient
+from airflow_mcp_server.parser.operation_parser import OperationParser
+from airflow_mcp_server.tools.airflow_tool import AirflowTool
+
+logger = logging.getLogger(__name__)
+
+_tools_cache: dict[str, AirflowTool] = {}
+
+
+def _initialize_client() -> AirflowClient:
+    """Initialize Airflow client with environment variables or embedded spec.
+
+    Returns:
+        AirflowClient instance
+
+    Raises:
+        ValueError: If required environment variables are missing or default spec is not found
+    """
+    spec_path = os.environ.get("OPENAPI_SPEC")
+    if not spec_path:
+        # Fallback to embedded v1.yaml
+        try:
+            with resources.files("airflow_mcp_server.resources").joinpath("v1.yaml").open("rb") as f:
+                spec_path = f.name
+                logger.info("OPENAPI_SPEC not set; using embedded v1.yaml from %s", spec_path)
+        except Exception as e:
+            raise ValueError("Default OpenAPI spec not found in package resources") from e
+
+    required_vars = ["AIRFLOW_BASE_URL", "AUTH_TOKEN"]
+    missing_vars = [var for var in required_vars if var not in os.environ]
+    if missing_vars:
+        raise ValueError(f"Missing required environment variables: {missing_vars}")
+
+    return AirflowClient(spec_path=spec_path, base_url=os.environ["AIRFLOW_BASE_URL"], auth_token=os.environ["AUTH_TOKEN"])
+
+
+async def _initialize_tools() -> None:
+    """Initialize tools cache with Airflow operations.
+
+    Raises:
+        ValueError: If initialization fails
+    """
+    global _tools_cache
+
+    try:
+        client = _initialize_client()
+        spec_path = os.environ.get("OPENAPI_SPEC")
+        if not spec_path:
+            with resources.files("airflow_mcp_server.resources").joinpath("v1.yaml").open("rb") as f:
+                spec_path = f.name
+        parser = OperationParser(spec_path)
+
+        # Generate tools for each operation
+        for operation_id in parser.get_operations():
+            operation_details = parser.parse_operation(operation_id)
+            tool = AirflowTool(operation_details, client)
+            _tools_cache[operation_id] = tool
+
+    except Exception as e:
+        logger.error("Failed to initialize tools: %s", e)
+        _tools_cache.clear()
+        raise ValueError(f"Failed to initialize tools: {e}") from e
+
+
+async def get_airflow_tools(mode: str = "unsafe") -> list[Tool]:
+    """Get list of available Airflow tools based on mode.
+
+    Args:
+        mode: "safe" for GET operations only, "unsafe" for all operations (default)
+
+    Returns:
+        List of MCP Tool objects representing available operations
+
+    Raises:
+        ValueError: If initialization fails
+    """
+    if not _tools_cache:
+        await _initialize_tools()
+
+    tools = []
+    for operation_id, tool in _tools_cache.items():
+        try:
+            # Skip non-GET operations in safe mode
+            if mode == "safe" and not tool.operation.method.lower() == "get":
+                continue
+            schema = tool.operation.input_model.model_json_schema()
+            tools.append(
+                Tool(
+                    name=operation_id,
+                    description=tool.operation.operation_id,
+                    inputSchema=schema,
+                )
+            )
+        except Exception as e:
+            logger.error("Failed to create tool schema for %s: %s", operation_id, e)
+            continue
+
+    return tools
+
+
+async def get_tool(name: str) -> AirflowTool:
+    """Get specific tool by name.
+
+    Args:
+        name: Tool/operation name
+
+    Returns:
+        AirflowTool instance
+
+    Raises:
+        KeyError: If tool not found
+        ValueError: If tool initialization fails
+    """
+    if not _tools_cache:
+        await _initialize_tools()
+
+    if name not in _tools_cache:
+        raise KeyError(f"Tool {name} not found")
+
+    return _tools_cache[name]