AirflowDAGMonitor/airflow_monitor/discovery.py

"""Auto-discovery of Airflow Docker infrastructure.

Finds running Airflow containers, locates docker-compose project directory,
reads .env and docker-compose.yml to extract connection parameters.
"""

import json
import logging
import os
import re
import subprocess
from pathlib import Path

import yaml

logger = logging.getLogger(__name__)


class DiscoveryError(Exception):
    """Raised when Airflow Docker infrastructure cannot be found."""


def discover_airflow() -> dict:
    """Auto-discover Airflow Docker setup and return connection parameters.

    Discovery steps:
    1. Find airflow-webserver container via `docker ps`
    2. Extract compose project directory from container labels
    3. Read .env from compose directory
    4. Read docker-compose.yml for port mappings and env vars
    5. Build connection config dict

    Returns:
        dict with keys: base_url, username, password, api_version,
        compose_dir, container_name
    """
    logger.info("Starting Airflow Docker auto-discovery")

    # Step 1: Find webserver container
    container = _find_webserver_container()
    container_name = container["Names"]
    logger.info("Found Airflow webserver container: %s", container_name)

    # Step 2: Get compose project directory
    compose_dir = _get_compose_dir(container)
    logger.info("Compose project directory: %s", compose_dir)

    # Step 3: Read .env file
    env_vars = _read_env_file(compose_dir)

    # Step 4: Read docker-compose.yml
    compose_config = _read_compose_file(compose_dir)

    # Step 5: Extract connection parameters
    host_port = _extract_webserver_port(container)
    credentials = _extract_credentials(env_vars, compose_config)

    result = {
        "base_url": f"http://localhost:{host_port}",
        "username": credentials["username"],
        "password": credentials["password"],
        "api_version": "v1",
        "compose_dir": str(compose_dir),
        "container_name": container_name,
    }

    logger.info(
        "Discovery complete: url=%s, user=%s, compose_dir=%s",
        result["base_url"], result["username"], result["compose_dir"],
    )
    return result


def _run_cmd(cmd: list[str], timeout: int = 15) -> str:
    """Run a shell command and return stdout. Raises DiscoveryError on failure."""
    logger.debug("Running command: %s", " ".join(cmd))
    try:
        result = subprocess.run(
            cmd, capture_output=True, text=True, timeout=timeout,
        )
        logger.debug(
            "Command result: rc=%d, stdout=%d bytes, stderr=%s",
            result.returncode, len(result.stdout),
            result.stderr.strip()[:200] if result.stderr.strip() else "<empty>",
        )
        if result.returncode != 0:
            raise DiscoveryError(
                f"Command failed: {' '.join(cmd)}\n"
                f"stderr: {result.stderr.strip()}"
            )
        return result.stdout.strip()
    except FileNotFoundError:
        raise DiscoveryError(
            f"Command not found: {cmd[0]}. Is Docker installed?"
        )
    except subprocess.TimeoutExpired:
        raise DiscoveryError(f"Command timed out: {' '.join(cmd)}")


def _find_webserver_container() -> dict:
    """Find running Airflow webserver container.

    Searches for containers with 'airflow' in the image and 'webserver'
    in the name or command.
    """
    output = _run_cmd([
        "docker", "ps", "--format", "{{json .}}",
        "--filter", "status=running",
    ])

    if not output:
        raise DiscoveryError("No running Docker containers found")

    candidates = []
    for line in output.splitlines():
        try:
            container = json.loads(line)
        except json.JSONDecodeError:
            continue

        image = container.get("Image", "").lower()
        names = container.get("Names", "").lower()
        command = container.get("Command", "").lower()
        status = container.get("Status", "")

        logger.debug(
            "  Container: name=%s, image=%s, command=%s, status=%s",
            names, image, command[:60], status,
        )

        # Match airflow webserver by multiple signals
        is_airflow = "airflow" in image or "airflow" in names
        is_webserver = (
            "webserver" in names
            or "webserver" in command
            or ("airflow" in command and "webserver" in command)
        )

        if is_airflow and is_webserver:
            logger.debug("  → MATCH: airflow webserver candidate")
            candidates.append(container)

    if not candidates:
        raise DiscoveryError(
            "No running Airflow webserver container found. "
            "Checked: image contains 'airflow' AND name/command contains 'webserver'"
        )

    if len(candidates) > 1:
        logger.warning(
            "Found %d webserver containers, using first: %s",
            len(candidates), candidates[0]["Names"],
        )

    return candidates[0]


def _get_compose_dir(container: dict) -> Path:
    """Get docker-compose project directory from container labels."""
    container_name = container["Names"]

    # Inspect container for compose labels
    output = _run_cmd([
        "docker", "inspect",
        "--format", '{{index .Config.Labels "com.docker.compose.project.working_dir"}}',
        container_name,
    ])

    if output and output != "<no value>":
        compose_dir = Path(output)
        if compose_dir.exists():
            return compose_dir

    # Fallback: try to find compose file via container's bind mounts
    inspect_json = _run_cmd(["docker", "inspect", container_name])
    try:
        data = json.loads(inspect_json)
        if data:
            mounts = data[0].get("Mounts", [])
            for mount in mounts:
                source = Path(mount.get("Source", ""))
                # Look for parent directory containing docker-compose.yml
                for candidate in [source.parent, source.parent.parent, source]:
                    for compose_name in ["docker-compose.yml", "docker-compose.yaml", "compose.yml", "compose.yaml"]:
                        if (candidate / compose_name).exists():
                            return candidate
    except (json.JSONDecodeError, IndexError, KeyError):
        pass

    raise DiscoveryError(
        f"Cannot determine compose directory for container {container_name}. "
        "Label 'com.docker.compose.project.working_dir' not found."
    )


def _extract_webserver_port(container: dict) -> int:
    """Extract host port mapped to webserver's 8080."""
    container_name = container["Names"]

    # docker port gives us the exact mapping
    try:
        output = _run_cmd(["docker", "port", container_name, "8080"])
        # Output: "0.0.0.0:80" or "0.0.0.0:80\n:::80"
        for line in output.splitlines():
            match = re.search(r":(\d+)$", line.strip())
            if match:
                port = int(match.group(1))
                logger.info("Webserver port: %d (from docker port)", port)
                return port
    except DiscoveryError:
        pass

    # Fallback: parse Ports field from docker ps
    ports_str = container.get("Ports", "")
    # Format: "0.0.0.0:80->8080/tcp"
    match = re.search(r"(\d+)->8080", ports_str)
    if match:
        port = int(match.group(1))
        logger.info("Webserver port: %d (from docker ps)", port)
        return port

    logger.warning("Cannot determine webserver port, defaulting to 8080")
    return 8080


def _read_env_file(compose_dir: Path) -> dict:
    """Read .env file from compose directory."""
    env_file = compose_dir / ".env"
    env_vars = {}

    if not env_file.exists():
        logger.warning("No .env file found in %s", compose_dir)
        return env_vars

    with open(env_file, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            if "=" in line:
                key, _, value = line.partition("=")
                key = key.strip()
                value = value.strip().strip("'\"")
                env_vars[key] = value

    logger.info("Read %d variables from .env", len(env_vars))
    # Log variable names (not values) for debugging
    logger.debug(
        ".env variables: %s",
        ", ".join(sorted(env_vars.keys())),
    )
    return env_vars


def _read_compose_file(compose_dir: Path) -> dict:
    """Read docker-compose.yml from compose directory."""
    for name in ["docker-compose.yml", "docker-compose.yaml", "compose.yml", "compose.yaml"]:
        path = compose_dir / name
        if path.exists():
            with open(path, "r", encoding="utf-8") as f:
                data = yaml.safe_load(f)
            logger.info("Read compose file: %s", path)
            return data if isinstance(data, dict) else {}

    logger.warning("No docker-compose.yml found in %s", compose_dir)
    return {}


def _extract_credentials(env_vars: dict, compose_config: dict) -> dict:
    """Extract Airflow webserver credentials from .env and compose config.

    Checks multiple variable names in priority order since different
    setups use different naming conventions.
    """
    # Username: check .env, then compose env, then default
    username = (
        env_vars.get("_AIRFLOW_WWW_USER_USERNAME")
        or env_vars.get("AIRFLOW_WWW_USER_USERNAME")
        or _get_compose_env(compose_config, "_AIRFLOW_WWW_USER_USERNAME")
        or "airflow"
    )

    # Password: check .env, then compose env, then default
    password = (
        env_vars.get("_AIRFLOW_WWW_USER_PASSWORD")
        or env_vars.get("AIRFLOW_WWW_USER_PASSWORD")
        or _get_compose_env(compose_config, "_AIRFLOW_WWW_USER_PASSWORD")
        or "airflow"
    )

    # Resolve ${VAR:-default} references in compose values
    password = _resolve_env_ref(password, env_vars)
    username = _resolve_env_ref(username, env_vars)

    logger.info("Credentials: user=%s, password=%s", username, "***")
    return {"username": username, "password": password}


def _get_compose_env(compose_config: dict, var_name: str) -> str | None:
    """Extract environment variable value from docker-compose services.

    Looks in airflow-init and airflow-webserver services.
    """
    services = compose_config.get("services", {})

    for service_name in ["airflow-init", "airflow-webserver"]:
        service = services.get(service_name, {})
        env = service.get("environment", {})

        if isinstance(env, dict):
            value = env.get(var_name)
            if value is not None:
                return str(value)
        elif isinstance(env, list):
            for item in env:
                if isinstance(item, str) and item.startswith(f"{var_name}="):
                    return item.split("=", 1)[1]

    return None


def _resolve_env_ref(value: str, env_vars: dict) -> str:
    """Resolve ${VAR:-default} or ${VAR} references in a string."""
    if not isinstance(value, str):
        return str(value) if value is not None else ""

    # Pattern: ${VAR_NAME:-default_value} or ${VAR_NAME}
    def replacer(match):
        var_name = match.group(1)
        default = match.group(3) if match.group(3) is not None else ""
        return env_vars.get(var_name, default)

    return re.sub(r"\$\{([^:}]+)(?::-(.*?))?\}", replacer, value)