astro-airflow-mcp

"""Main entry point for running the Airflow MCP server."""

import argparse
import logging
import os
from pathlib import Path

import yaml

from astro_airflow_mcp.logging import configure_logging, get_logger
from astro_airflow_mcp.server import configure, mcp

logger = get_logger("main")

# Default Airflow URL if no config is found
DEFAULT_AIRFLOW_URL = "http://localhost:8080"


def discover_airflow_url(project_dir: str | None) -> str | None:
    """Discover Airflow URL from .astro/config.yaml in the project directory.

    Looks for the Astro CLI config file and extracts the webserver/api-server port.
    Prefers api-server.port (Airflow 3.x) over webserver.port (Airflow 2.x).

    Args:
        project_dir: The project directory to search in

    Returns:
        The discovered Airflow URL (e.g., "http://localhost:8081"), or None if not found
    """
    if not project_dir:
        return None

    config_path = Path(project_dir) / ".astro" / "config.yaml"
    if not config_path.exists():
        return None

    try:
        with open(config_path) as f:
            config = yaml.safe_load(f)

        if not config:
            return None

        # Try api-server.port first (Airflow 3.x), then webserver.port (Airflow 2.x)
        port = None
        if "api-server" in config and isinstance(config["api-server"], dict):
            port = config["api-server"].get("port")
        if port is None and "webserver" in config and isinstance(config["webserver"], dict):
            port = config["webserver"].get("port")

        if port:
            return f"http://localhost:{port}"

    except Exception as e:
        # Log but don't fail - we'll fall back to default
        logger.debug("Failed to read .astro/config.yaml: %s", e)

    return None


def main():
    """Main entry point for the Airflow MCP server."""
    # Parse command line arguments first to determine transport mode
    parser = argparse.ArgumentParser(description="Airflow MCP Server")
    parser.add_argument(
        "--transport",
        type=str,
        default=os.getenv("MCP_TRANSPORT", "stdio"),
        choices=["stdio", "http"],
        help="Transport mode: stdio (default) or http",
    )
    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"),  # None if not set
        help="Base URL of Airflow webserver (auto-discovered from .astro/config.yaml if not provided)",
    )
    parser.add_argument(
        "--auth-token",
        type=str,
        default=os.getenv("AIRFLOW_AUTH_TOKEN"),
        help="Bearer token for Airflow API authentication (takes precedence over username/password)",
    )
    parser.add_argument(
        "--username",
        type=str,
        default=os.getenv("AIRFLOW_USERNAME"),
        help="Username for Airflow API token authentication",
    )
    parser.add_argument(
        "--password",
        type=str,
        default=os.getenv("AIRFLOW_PASSWORD"),
        help="Password for Airflow API token authentication",
    )
    parser.add_argument(
        "--airflow-project-dir",
        type=str,
        default=os.getenv("AIRFLOW_PROJECT_DIR") or os.getenv("PWD") or os.getcwd(),
        help="Astro project directory for auto-discovering Airflow URL from .astro/config.yaml (default: $PWD)",
    )

    args = parser.parse_args()

    # Configure logging - use stderr in stdio mode to avoid corrupting JSON-RPC
    stdio_mode = args.transport == "stdio"
    configure_logging(level=logging.INFO, stdio_mode=stdio_mode)

    # Determine Airflow URL: explicit > auto-discover > default
    airflow_url = args.airflow_url
    url_source = "explicit"
    if not airflow_url:
        # Try auto-discovery from .astro/config.yaml
        airflow_url = discover_airflow_url(args.airflow_project_dir)
        if airflow_url:
            url_source = "auto-discovered"
        else:
            airflow_url = DEFAULT_AIRFLOW_URL
            url_source = "default"

    # Configure Airflow connection settings
    configure(
        url=airflow_url,
        auth_token=args.auth_token,
        username=args.username,
        password=args.password,
        project_dir=args.airflow_project_dir,
    )

    # Log configuration
    logger.info("Project directory: %s", args.airflow_project_dir)
    logger.info("Airflow URL: %s (%s)", airflow_url, url_source)
    if args.auth_token:
        logger.info("Authentication: Direct bearer token")
    elif args.username:
        logger.info("Authentication: Token manager (username: %s)", args.username)
    else:
        logger.info("Authentication: Token manager (credential-less mode)")

    # 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()