diff --git a/airflow-mcp-server/pyproject.toml b/airflow-mcp-server/pyproject.toml index 442016e..671a747 100644 --- a/airflow-mcp-server/pyproject.toml +++ b/airflow-mcp-server/pyproject.toml @@ -38,6 +38,14 @@ exclude = [ "!pyproject.toml" ] +[tool.hatch.build.targets.wheel] +packages = ["src/airflow_mcp_server"] + +[tool.hatch.build.targets.wheel.sources] +"src/airflow_mcp_server" = [ + "*.yaml", +] + [tool.pytest.ini_options] pythonpath = ["src"] asyncio_mode = "strict" diff --git a/airflow-mcp-server/src/airflow_mcp_server/server.py b/airflow-mcp-server/src/airflow_mcp_server/server.py index eeb2ca1..09933d2 100644 --- a/airflow-mcp-server/src/airflow_mcp_server/server.py +++ b/airflow-mcp-server/src/airflow_mcp_server/server.py @@ -22,7 +22,7 @@ logger = logging.getLogger(__name__) async def serve() -> None: """Start MCP server.""" - required_vars = ["OPENAPI_SPEC", "AIRFLOW_BASE_URL", "AUTH_TOKEN"] + required_vars = ["AIRFLOW_BASE_URL", "AUTH_TOKEN"] if not all(var in os.environ for var in required_vars): raise ValueError(f"Missing required environment variables: {required_vars}") diff --git a/airflow-mcp-server/src/airflow_mcp_server/tools/tool_manager.py b/airflow-mcp-server/src/airflow_mcp_server/tools/tool_manager.py index 96af499..74df2e5 100644 --- a/airflow-mcp-server/src/airflow_mcp_server/tools/tool_manager.py +++ b/airflow-mcp-server/src/airflow_mcp_server/tools/tool_manager.py @@ -1,5 +1,6 @@ import logging import os +from importlib import resources from mcp.types import Tool @@ -13,20 +14,30 @@ _tools_cache: dict[str, AirflowTool] = {} def _initialize_client() -> AirflowClient: - """Initialize Airflow client with environment variables. + """Initialize Airflow client with environment variables or embedded spec. Returns: AirflowClient instance Raises: - ValueError: If required environment variables are missing + ValueError: If required environment variables are missing or default spec is not found """ - required_vars = ["OPENAPI_SPEC", "AIRFLOW_BASE_URL", "AUTH_TOKEN"] + 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=os.environ["OPENAPI_SPEC"], base_url=os.environ["AIRFLOW_BASE_URL"], auth_token=os.environ["AUTH_TOKEN"]) + return AirflowClient(spec_path=spec_path, base_url=os.environ["AIRFLOW_BASE_URL"], auth_token=os.environ["AUTH_TOKEN"]) async def _initialize_tools() -> None: @@ -39,7 +50,11 @@ async def _initialize_tools() -> None: try: client = _initialize_client() - parser = OperationParser(os.environ["OPENAPI_SPEC"]) + 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(): @@ -60,7 +75,7 @@ async def get_airflow_tools() -> list[Tool]: List of MCP Tool objects representing available operations Raises: - ValueError: If required environment variables are missing or initialization fails + ValueError: If initialization fails """ if not _tools_cache: await _initialize_tools()