"""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 "", ) 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 != "": 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)