AirflowDAGMonitor/airflow_monitor/__main__.py

"""Entry point for Airflow DAG Monitor.

Usage:
    python -m airflow_monitor --config /path/to/config.yaml
    python -m airflow_monitor --discover              # auto-detect Docker Airflow
    python -m airflow_monitor --discover --dry-run     # show discovered config
"""

import argparse
import dataclasses
import fcntl
import json
import logging
import logging.handlers
import os
import pathlib
import signal
import sys
import threading

from .config import AppConfig, load_config
from .discovery import DiscoveryError, discover_airflow
from .monitor import Monitor


def setup_logging(config) -> None:
    """Configure logging with rotating file handler and stdout."""
    root_logger = logging.getLogger()
    root_logger.setLevel(getattr(logging, config.level.upper(), logging.INFO))

    formatter = logging.Formatter(
        "%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )

    # Stdout handler (captured by journald when running as service)
    stdout_handler = logging.StreamHandler(sys.stdout)
    stdout_handler.setFormatter(formatter)
    root_logger.addHandler(stdout_handler)

    # Rotating file handler
    log_path = pathlib.Path(config.file)
    log_path.parent.mkdir(parents=True, exist_ok=True)
    try:
        file_handler = logging.handlers.RotatingFileHandler(
            config.file,
            maxBytes=config.max_bytes,
            backupCount=config.backup_count,
            encoding="utf-8",
        )
        file_handler.setFormatter(formatter)
        root_logger.addHandler(file_handler)
    except OSError as e:
        root_logger.warning("Cannot open log file %s: %s (using stdout only)", config.file, e)


def acquire_lock(lock_path: str):
    """Acquire an exclusive lock to prevent multiple instances.

    Returns the file descriptor (must be kept open for lock duration).
    Exits with code 1 if another instance is running.
    """
    lock_dir = pathlib.Path(lock_path).parent
    lock_dir.mkdir(parents=True, exist_ok=True)

    try:
        fd = open(lock_path, "w")
        fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
        fd.write(str(os.getpid()))
        fd.flush()
        return fd
    except BlockingIOError:
        print(
            f"Another instance is already running (lock: {lock_path})",
            file=sys.stderr,
        )
        sys.exit(1)
    except OSError as e:
        print(f"Cannot acquire lock file {lock_path}: {e}", file=sys.stderr)
        sys.exit(1)


def _apply_discovery(config: AppConfig, discovered: dict) -> AppConfig:
    """Override airflow connection settings with discovered values."""
    airflow = dataclasses.replace(
        config.airflow,
        base_url=discovered["base_url"],
        username=discovered["username"],
        password=discovered["password"],
        api_version=discovered.get("api_version", "v1"),
        verify_ssl=False,  # Docker HTTP, not HTTPS
    )
    return dataclasses.replace(config, airflow=airflow)


def main():
    parser = argparse.ArgumentParser(
        description="Airflow DAG Monitor - monitors DAG runs and alerts via Zabbix",
    )
    parser.add_argument(
        "--config", "-c",
        default="config.yaml",
        help="Path to config.yaml (default: config.yaml in current directory)",
    )
    parser.add_argument(
        "--discover", "-d",
        action="store_true",
        help="Auto-discover Airflow Docker setup (overrides airflow section in config)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="With --discover: show discovered config and exit without starting monitor",
    )
    args = parser.parse_args()

    # Load config
    try:
        config = load_config(args.config)
    except (FileNotFoundError, ValueError) as e:
        if args.discover:
            # Config file optional when using discovery — use defaults
            from .config import AirflowConfig, MonitorConfig, ZabbixConfig, LoggingConfig
            config = AppConfig(
                airflow=AirflowConfig(),
                monitor=MonitorConfig(),
                zabbix=ZabbixConfig(),
                logging=LoggingConfig(),
            )
        else:
            print(f"Configuration error: {e}", file=sys.stderr)
            sys.exit(1)

    # Auto-discovery
    if args.discover:
        try:
            discovered = discover_airflow()
            config = _apply_discovery(config, discovered)
            print(f"Discovered Airflow at {discovered['base_url']}")
            print(f"  Compose dir: {discovered['compose_dir']}")
            print(f"  Container:   {discovered['container_name']}")
            print(f"  User:        {discovered['username']}")
        except DiscoveryError as e:
            print(f"Discovery failed: {e}", file=sys.stderr)
            sys.exit(1)

        if args.dry_run:
            print("\n--- Effective Airflow config ---")
            print(json.dumps(dataclasses.asdict(config.airflow), indent=2))
            print("\n--- Effective Zabbix config ---")
            print(json.dumps(dataclasses.asdict(config.zabbix), indent=2))
            sys.exit(0)

    # Setup logging
    setup_logging(config.logging)
    logger = logging.getLogger(__name__)
    logger.info("Airflow DAG Monitor starting")

    # Acquire lock
    lock_fd = acquire_lock(config.monitor.lock_file)
    logger.info("Lock acquired: %s (PID %d)", config.monitor.lock_file, os.getpid())

    # Setup graceful shutdown
    shutdown_event = threading.Event()

    def handle_signal(signum, frame):
        sig_name = signal.Signals(signum).name
        logger.info("Received %s, initiating shutdown", sig_name)
        shutdown_event.set()

    signal.signal(signal.SIGTERM, handle_signal)
    signal.signal(signal.SIGINT, handle_signal)

    # Run monitor
    try:
        monitor = Monitor(config, shutdown_event)
        monitor.run()
    except Exception:
        logger.exception("Fatal error")
        sys.exit(1)
    finally:
        lock_fd.close()
        logger.info("Monitor stopped")


if __name__ == "__main__":
    main()