astro-airflow-mcp 0.1.5__py3-none-any.whl

astro_airflow_mcp/__init__.py

@@ -0,0 +1,5 @@
+"""Airflow MCP Server."""
+
+from importlib.metadata import version
+
+__version__ = version("astro-airflow-mcp")

astro_airflow_mcp/__main__.py

@@ -0,0 +1,75 @@
+"""Main entry point for running the Airflow MCP server."""
+
+import argparse
+import logging
+import os
+
+from astro_airflow_mcp.logging import configure_logging, get_logger
+from astro_airflow_mcp.server import configure, mcp
+
+logger = get_logger("main")
+
+
+def main():
+    """Main entry point for the Airflow MCP server."""
+    # Configure logging
+    configure_logging(level=logging.INFO)
+
+    # Parse command line arguments
+    parser = argparse.ArgumentParser(description="Airflow MCP Server")
+    parser.add_argument(
+        "--transport",
+        type=str,
+        default=os.getenv("MCP_TRANSPORT", "http"),
+        choices=["stdio", "http"],
+        help="Transport mode: http (default) or stdio",
+    )
+    parser.add_argument(
+        "--host",
+        type=str,
+        default=os.getenv("MCP_HOST", "localhost"),
+        help="Host to bind to (only for http transport, default: localhost)",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=int(os.getenv("MCP_PORT", "8000")),
+        help="Port to bind to (only for http transport, default: 8000)",
+    )
+    parser.add_argument(
+        "--airflow-url",
+        type=str,
+        default=os.getenv("AIRFLOW_API_URL", "http://localhost:8080"),
+        help="Base URL of Airflow webserver (default: http://localhost:8080)",
+    )
+    parser.add_argument(
+        "--auth-token",
+        type=str,
+        default=os.getenv("AIRFLOW_AUTH_TOKEN"),
+        help="Bearer token for Airflow API authentication",
+    )
+
+    args = parser.parse_args()
+
+    # Configure Airflow connection settings
+    configure(
+        url=args.airflow_url,
+        auth_token=args.auth_token,
+    )
+
+    # Log Airflow connection configuration
+    logger.info(f"Airflow URL: {args.airflow_url}")
+    if args.auth_token:
+        logger.info("Authentication: Bearer token")
+    else:
+        logger.info("Authentication: None")
+
+    # Run the server with specified transport
+    if args.transport == "http":
+        mcp.run(transport="http", host=args.host, port=args.port, show_banner=False)
+    else:
+        mcp.run(show_banner=False)
+
+
+if __name__ == "__main__":
+    main()

astro_airflow_mcp/logging.py

@@ -0,0 +1,61 @@
+"""Logging utilities for Airflow MCP."""
+
+import logging
+import sys
+
+
+def get_logger(name: str | None = None) -> logging.Logger:
+    """Get a logger instance nested under the airflow_mcp namespace.
+
+    Args:
+        name: Optional name for the logger. If not provided, returns root airflow_mcp logger.
+              Will be nested under 'airflow_mcp' (e.g., 'airflow_mcp.server')
+
+    Returns:
+        A logger instance
+
+    Example:
+        >>> logger = get_logger("server")
+        >>> logger.info("Starting server")
+    """
+    if name:
+        return logging.getLogger(f"airflow_mcp.{name}")
+    return logging.getLogger("airflow_mcp")
+
+
+def configure_logging(level: str | int = logging.INFO) -> None:
+    """Configure logging for the airflow_mcp package.
+
+    Sets up a simple console handler with a standard format.
+
+    Args:
+        level: Logging level (e.g., logging.INFO, logging.DEBUG, or "INFO", "DEBUG")
+              Defaults to INFO.
+
+    Example:
+        >>> configure_logging(level=logging.DEBUG)
+    """
+    # Convert string level to int if needed
+    if isinstance(level, str):
+        level = getattr(logging, level.upper(), logging.INFO)
+
+    # Get the root airflow_mcp logger
+    logger = get_logger()
+    logger.setLevel(level)
+
+    # Remove any existing handlers to avoid duplicates
+    logger.handlers.clear()
+
+    # Create console handler
+    handler = logging.StreamHandler(sys.stdout)
+    handler.setLevel(level)
+
+    # Create simple formatter
+    formatter = logging.Formatter(fmt="%(levelname)s: %(message)s")
+    handler.setFormatter(formatter)
+
+    # Add handler to logger
+    logger.addHandler(handler)
+
+    # Prevent propagation to root logger to avoid duplicate logs
+    logger.propagate = False

astro_airflow_mcp/plugin.py

@@ -0,0 +1,66 @@
+"""Airflow plugin for integrating MCP server."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from astro_airflow_mcp import __version__
+
+# Use standard logging for Airflow plugin integration
+# This allows Airflow to control log level, format, and destination
+logger = logging.getLogger(__name__)
+
+try:
+    from airflow.plugins_manager import AirflowPlugin
+
+    AIRFLOW_AVAILABLE = True
+except ImportError:
+    AIRFLOW_AVAILABLE = False
+    AirflowPlugin = object  # type: ignore
+    logger.warning("Airflow not available, plugin disabled")
+
+
+# FastAPI app configuration for Airflow 3.x plugin system
+try:
+    from fastapi import FastAPI
+
+    from astro_airflow_mcp.server import mcp
+
+    # Get the native MCP protocol ASGI app from FastMCP
+    mcp_protocol_app = mcp.http_app(path="/")
+
+    # Wrap in a FastAPI app with the MCP app's lifespan
+    # This is required for FastMCP to initialize its task group
+    app = FastAPI(
+        title="Airflow MCP Server", version=__version__, lifespan=mcp_protocol_app.lifespan
+    )
+
+    # Mount the MCP protocol app
+    app.mount("/v1", mcp_protocol_app)
+    logger.info("MCP protocol app created and mounted")
+
+    # Airflow plugin configuration
+    fastapi_apps_config = [{"app": app, "url_prefix": "/mcp", "name": "Airflow MCP Server"}]
+
+except ImportError as e:
+    logger.warning(f"FastAPI integration not available: {e}")
+    fastapi_apps_config = []
+
+
+class AirflowMCPPlugin(AirflowPlugin):
+    """Plugin to integrate MCP server with Airflow.
+
+    Exposes MCP protocol endpoints at /mcp for AI clients (Cursor, Claude Desktop, etc.)
+    """
+
+    name = "astro_airflow_mcp"
+    fastapi_apps = fastapi_apps_config
+
+    @staticmethod
+    def on_load(*_args: Any, **_kwargs: Any) -> None:
+        """Called when the plugin is loaded."""
+        logger.info("Airflow MCP Plugin loaded")
+
+
+__all__ = ["AirflowMCPPlugin"]