feat: добавил систему мониторинга Airflow DAG-ов с интеграцией в Zabbix
- автообнаружение Docker Airflow через docker inspect и .env - мониторинг failed и long-running DAG-запусков с автоматическим retry - экспорт данных в файлы для Zabbix Agent через UserParameter - офлайн-сборка ZIP-архива для закрытых контуров - Zabbix шаблоны для 5.x (XML) и 6.x+ (YAML) - systemd сервис с graceful shutdown и lock file
This commit is contained in:
Maksim Totmin
@@ -0,0 +1,3 @@
|
||||
"""Airflow DAG Monitor - monitors DAG runs and alerts via Zabbix."""
|
||||
|
||||
__version__ = "1.0.0"
|
||||
@@ -0,0 +1,187 @@
|
||||
"""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()
|
||||
@@ -0,0 +1,244 @@
|
||||
"""Action handlers: retry DAG runs and export data for Zabbix agent."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import pathlib
|
||||
import tempfile
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from .analyzer import DagIssue
|
||||
from .client import AirflowClient
|
||||
from .config import MonitorConfig, ZabbixConfig
|
||||
from .state import StateManager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DataExporter:
|
||||
"""Exports monitoring data to files for Zabbix agent UserParameter.
|
||||
|
||||
Files written:
|
||||
- problems.json — JSON array of problematic DAG runs (or [])
|
||||
- heartbeat — epoch timestamp (updated every cycle)
|
||||
- status.json — overall monitor status (cycle count, DAG count, etc.)
|
||||
|
||||
Zabbix agent reads these via UserParameter defined in
|
||||
/etc/zabbix/zabbix_agentd.conf.d/airflow-monitor.conf
|
||||
"""
|
||||
|
||||
def __init__(self, config: ZabbixConfig):
|
||||
self._config = config
|
||||
self._data_dir = pathlib.Path(config.data_dir)
|
||||
self._problems_path = self._data_dir / config.problems_file
|
||||
self._heartbeat_path = self._data_dir / config.heartbeat_file
|
||||
self._status_path = self._data_dir / config.status_file
|
||||
|
||||
logger.debug(
|
||||
"DataExporter initialized: enabled=%s, data_dir=%s, "
|
||||
"problems=%s, heartbeat=%s, status=%s",
|
||||
config.enabled, self._data_dir,
|
||||
self._problems_path, self._heartbeat_path, self._status_path,
|
||||
)
|
||||
|
||||
if config.enabled:
|
||||
self._data_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def _atomic_write(self, path: pathlib.Path, content: str):
|
||||
"""Write content to file atomically (write tmp → rename)."""
|
||||
try:
|
||||
fd, tmp_path = tempfile.mkstemp(
|
||||
dir=str(path.parent), suffix=".tmp",
|
||||
)
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
||||
f.write(content)
|
||||
os.replace(tmp_path, str(path))
|
||||
logger.debug(
|
||||
"File written: %s (%d bytes)",
|
||||
path, len(content),
|
||||
)
|
||||
except OSError as e:
|
||||
logger.error("Failed to write %s: %s", path, e)
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def export_problems(self, problems: list[dict]) -> bool:
|
||||
"""Write problem list as JSON file for Zabbix agent.
|
||||
|
||||
Writes '[]' if no problems (Zabbix trigger auto-clears).
|
||||
"""
|
||||
if not self._config.enabled:
|
||||
logger.debug("DataExporter disabled, skipping export")
|
||||
return True
|
||||
|
||||
payload = json.dumps(problems, indent=2, ensure_ascii=False)
|
||||
logger.debug(
|
||||
"Exporting problems: count=%d, size=%d bytes → %s",
|
||||
len(problems), len(payload), self._problems_path,
|
||||
)
|
||||
self._atomic_write(self._problems_path, payload)
|
||||
return True
|
||||
|
||||
def export_heartbeat(self):
|
||||
"""Write current epoch timestamp for Zabbix agent heartbeat check."""
|
||||
if not self._config.enabled:
|
||||
return
|
||||
|
||||
epoch = str(int(time.time()))
|
||||
logger.debug("Exporting heartbeat: %s → %s", epoch, self._heartbeat_path)
|
||||
self._atomic_write(self._heartbeat_path, epoch)
|
||||
|
||||
def export_status(self, cycle_count: int, dag_count: int,
|
||||
issue_count: int, cycle_time: float):
|
||||
"""Write overall monitor status for Zabbix agent."""
|
||||
if not self._config.enabled:
|
||||
return
|
||||
|
||||
status = {
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"cycle_count": cycle_count,
|
||||
"dag_count": dag_count,
|
||||
"issue_count": issue_count,
|
||||
"cycle_time_seconds": round(cycle_time, 2),
|
||||
}
|
||||
payload = json.dumps(status, indent=2, ensure_ascii=False)
|
||||
logger.debug(
|
||||
"Exporting status: cycle=#%d, dags=%d, issues=%d, time=%.1fs → %s",
|
||||
cycle_count, dag_count, issue_count, cycle_time, self._status_path,
|
||||
)
|
||||
self._atomic_write(self._status_path, payload)
|
||||
|
||||
|
||||
class ActionHandler:
|
||||
"""Handles issue resolution: retry or escalate via data export."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: AirflowClient,
|
||||
state: StateManager,
|
||||
exporter: DataExporter,
|
||||
config: MonitorConfig,
|
||||
):
|
||||
self._client = client
|
||||
self._state = state
|
||||
self._exporter = exporter
|
||||
self._config = config
|
||||
logger.debug(
|
||||
"ActionHandler initialized: max_retries=%d, retry_wait=%ds",
|
||||
config.max_retries, config.retry_wait,
|
||||
)
|
||||
|
||||
def handle_issue(self, issue: DagIssue) -> str:
|
||||
"""Process a single issue.
|
||||
|
||||
Returns:
|
||||
"already_handled" - previously alerted, skip
|
||||
"retried" - retry initiated, needs re-check
|
||||
"needs_alert" - retry exhausted, alert needed
|
||||
"""
|
||||
dag_id = issue.dag_id
|
||||
run_id = issue.dag_run_id
|
||||
|
||||
# Ensure entry exists in state
|
||||
self._state.ensure_entry(dag_id, run_id)
|
||||
|
||||
retry_count = self._state.get_retry_count(dag_id, run_id)
|
||||
is_alerted = self._state.is_alerted(dag_id, run_id)
|
||||
|
||||
logger.debug(
|
||||
"Handling issue: %s/%s type=%s state=%s duration=%.0fs "
|
||||
"retry_count=%d/%d alerted=%s",
|
||||
dag_id, run_id, issue.issue_type, issue.state,
|
||||
issue.duration_seconds, retry_count, self._config.max_retries,
|
||||
is_alerted,
|
||||
)
|
||||
|
||||
# Already alerted? Skip.
|
||||
if is_alerted:
|
||||
logger.debug(" → already_handled: alert was sent previously")
|
||||
return "already_handled"
|
||||
|
||||
# Can we retry?
|
||||
if retry_count < self._config.max_retries:
|
||||
logger.debug(
|
||||
" → attempting retry %d/%d via clear_dag_run",
|
||||
retry_count + 1, self._config.max_retries,
|
||||
)
|
||||
success = self._client.clear_dag_run(dag_id, run_id)
|
||||
if success:
|
||||
self._state.increment_retry(dag_id, run_id)
|
||||
logger.info(
|
||||
"Retry %d/%d initiated for %s/%s (%s, duration=%.0fs)",
|
||||
retry_count + 1, self._config.max_retries,
|
||||
dag_id, run_id, issue.issue_type, issue.duration_seconds,
|
||||
)
|
||||
return "retried"
|
||||
else:
|
||||
logger.warning(
|
||||
"Retry failed for %s/%s (API error), escalating to alert",
|
||||
dag_id, run_id,
|
||||
)
|
||||
return "needs_alert"
|
||||
|
||||
# Retries exhausted
|
||||
logger.debug(
|
||||
" → needs_alert: retries exhausted (%d/%d)",
|
||||
retry_count, self._config.max_retries,
|
||||
)
|
||||
return "needs_alert"
|
||||
|
||||
def collect_and_alert(self, issues: list[DagIssue]):
|
||||
"""Export alerts for issues that need alerting.
|
||||
|
||||
Filters to only unalerted issues, builds JSON payload,
|
||||
writes to file for Zabbix agent, and marks as alerted in state.
|
||||
"""
|
||||
to_alert = []
|
||||
for issue in issues:
|
||||
if not self._state.is_alerted(issue.dag_id, issue.dag_run_id):
|
||||
entry = self._state.ensure_entry(issue.dag_id, issue.dag_run_id)
|
||||
alert_data = {
|
||||
"dag_id": issue.dag_id,
|
||||
"dag_run_id": issue.dag_run_id,
|
||||
"issue_type": issue.issue_type,
|
||||
"status": issue.state,
|
||||
"duration_seconds": issue.duration_seconds,
|
||||
"error_info": issue.error_info,
|
||||
"retry_count": entry["retry_count"],
|
||||
}
|
||||
to_alert.append(alert_data)
|
||||
logger.debug(
|
||||
"Issue queued for alert: %s/%s type=%s retries=%d",
|
||||
issue.dag_id, issue.dag_run_id,
|
||||
issue.issue_type, entry["retry_count"],
|
||||
)
|
||||
|
||||
if not to_alert:
|
||||
# Write empty list to clear Zabbix trigger
|
||||
self._exporter.export_problems([])
|
||||
logger.info("No issues to alert, exported empty problem list")
|
||||
return
|
||||
|
||||
logger.warning(
|
||||
"Exporting %d problematic DAG runs for Zabbix: %s",
|
||||
len(to_alert),
|
||||
", ".join(
|
||||
f"{p['dag_id']}/{p['dag_run_id']}({p['issue_type']})"
|
||||
for p in to_alert
|
||||
),
|
||||
)
|
||||
|
||||
# Log full alert payload at debug level
|
||||
logger.debug("Alert payload:\n%s", json.dumps(to_alert, indent=2, ensure_ascii=False))
|
||||
|
||||
self._exporter.export_problems(to_alert)
|
||||
|
||||
for issue in issues:
|
||||
if not self._state.is_alerted(issue.dag_id, issue.dag_run_id):
|
||||
self._state.mark_alerted(issue.dag_id, issue.dag_run_id)
|
||||
logger.debug("Marked as alerted: %s/%s", issue.dag_id, issue.dag_run_id)
|
||||
|
||||
logger.info("Exported %d issues to problems file for Zabbix agent", len(to_alert))
|
||||
@@ -0,0 +1,180 @@
|
||||
"""DAG run analyzer - classifies runs and detects issues."""
|
||||
|
||||
import dataclasses
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
from typing import Callable
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class DagIssue:
|
||||
"""Represents a detected problem with a DAG run."""
|
||||
|
||||
dag_id: str
|
||||
dag_run_id: str
|
||||
issue_type: str # "failed" or "long_running"
|
||||
state: str # airflow state string
|
||||
start_time: str # ISO 8601 UTC
|
||||
duration_seconds: float
|
||||
error_info: str = ""
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return dataclasses.asdict(self)
|
||||
|
||||
|
||||
class DagAnalyzer:
|
||||
"""Analyzes DAG runs and produces a list of issues."""
|
||||
|
||||
def __init__(self, long_running_threshold: int):
|
||||
self._threshold = long_running_threshold
|
||||
logger.debug(
|
||||
"DagAnalyzer initialized: long_running_threshold=%ds (%.1f min)",
|
||||
long_running_threshold, long_running_threshold / 60,
|
||||
)
|
||||
|
||||
def analyze_dag_runs(
|
||||
self,
|
||||
dag_id: str,
|
||||
runs: list[dict],
|
||||
task_instances_fn: Callable[[str, str], list[dict]],
|
||||
) -> list[DagIssue]:
|
||||
"""Analyze DAG runs for a single DAG.
|
||||
|
||||
Args:
|
||||
dag_id: DAG identifier.
|
||||
runs: List of DAG run dicts from Airflow API.
|
||||
task_instances_fn: Callable(dag_id, dag_run_id) -> list of task instances.
|
||||
Called lazily only for failed runs to extract error details.
|
||||
|
||||
Returns:
|
||||
List of DagIssue objects for problematic runs.
|
||||
"""
|
||||
issues = []
|
||||
now = datetime.now(timezone.utc)
|
||||
logger.debug("Analyzing %d runs for DAG '%s' (now=%s)", len(runs), dag_id, now.isoformat())
|
||||
|
||||
for run in runs:
|
||||
run_id = run.get("dag_run_id", "")
|
||||
state = run.get("state", "")
|
||||
start_date_str = run.get("start_date") or run.get("execution_date", "")
|
||||
|
||||
if not start_date_str:
|
||||
logger.warning("DAG run %s/%s has no start_date, skipping", dag_id, run_id)
|
||||
continue
|
||||
|
||||
start_date = self._parse_datetime(start_date_str)
|
||||
if start_date is None:
|
||||
logger.warning(
|
||||
"Cannot parse start_date '%s' for %s/%s",
|
||||
start_date_str, dag_id, run_id,
|
||||
)
|
||||
continue
|
||||
|
||||
duration = (now - start_date).total_seconds()
|
||||
|
||||
logger.debug(
|
||||
" Run %s/%s: state=%s, start=%s, duration=%.0fs (%.1f min), threshold=%ds",
|
||||
dag_id, run_id, state, start_date_str,
|
||||
duration, duration / 60, self._threshold,
|
||||
)
|
||||
|
||||
if state == "failed":
|
||||
error_info = self._extract_error(dag_id, run_id, task_instances_fn)
|
||||
issue = DagIssue(
|
||||
dag_id=dag_id,
|
||||
dag_run_id=run_id,
|
||||
issue_type="failed",
|
||||
state=state,
|
||||
start_time=start_date.isoformat(),
|
||||
duration_seconds=round(duration, 1),
|
||||
error_info=error_info,
|
||||
)
|
||||
issues.append(issue)
|
||||
logger.debug(" → ISSUE DETECTED: %s (error: %s)", issue.issue_type, error_info or "n/a")
|
||||
|
||||
elif state == "running" and duration > self._threshold:
|
||||
issue = DagIssue(
|
||||
dag_id=dag_id,
|
||||
dag_run_id=run_id,
|
||||
issue_type="long_running",
|
||||
state=state,
|
||||
start_time=start_date.isoformat(),
|
||||
duration_seconds=round(duration, 1),
|
||||
)
|
||||
issues.append(issue)
|
||||
logger.debug(
|
||||
" → ISSUE DETECTED: long_running (%.0fs > %ds threshold)",
|
||||
duration, self._threshold,
|
||||
)
|
||||
|
||||
elif state == "running":
|
||||
logger.debug(" → OK: running within threshold (%.0fs <= %ds)", duration, self._threshold)
|
||||
|
||||
else:
|
||||
logger.debug(" → SKIP: state=%s (not actionable)", state)
|
||||
|
||||
logger.debug("Analysis complete for DAG '%s': %d issues found", dag_id, len(issues))
|
||||
return issues
|
||||
|
||||
def _extract_error(
|
||||
self,
|
||||
dag_id: str,
|
||||
dag_run_id: str,
|
||||
task_instances_fn: Callable[[str, str], list[dict]],
|
||||
) -> str:
|
||||
"""Extract error info from failed task instances."""
|
||||
logger.debug("Extracting error info for %s/%s", dag_id, dag_run_id)
|
||||
try:
|
||||
tasks = task_instances_fn(dag_id, dag_run_id)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to fetch task instances for %s/%s: %s", dag_id, dag_run_id, e)
|
||||
return ""
|
||||
|
||||
failed_tasks = [t for t in tasks if t.get("state") == "failed"]
|
||||
logger.debug(
|
||||
"Task instances for %s/%s: total=%d, failed=%d",
|
||||
dag_id, dag_run_id, len(tasks), len(failed_tasks),
|
||||
)
|
||||
|
||||
if not failed_tasks:
|
||||
return ""
|
||||
|
||||
# Return info about the first failed task
|
||||
task = failed_tasks[0]
|
||||
task_id = task.get("task_id", "unknown")
|
||||
# Try to get the error message from different possible fields
|
||||
error = (
|
||||
task.get("rendered_fields", {}).get("error", "")
|
||||
or task.get("note", "")
|
||||
or ""
|
||||
)
|
||||
# Truncate long error messages
|
||||
if len(error) > 500:
|
||||
error = error[:500] + "..."
|
||||
|
||||
result = f"Task '{task_id}' failed" + (f": {error}" if error else "")
|
||||
logger.debug("Error extracted for %s/%s: %s", dag_id, dag_run_id, result)
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _parse_datetime(dt_string: str) -> datetime | None:
|
||||
"""Parse ISO 8601 datetime string from Airflow API.
|
||||
|
||||
Handles both 'Z' suffix and '+00:00' timezone offset.
|
||||
"""
|
||||
if not dt_string:
|
||||
return None
|
||||
|
||||
# Normalize 'Z' to '+00:00' for fromisoformat
|
||||
cleaned = dt_string.replace("Z", "+00:00")
|
||||
|
||||
try:
|
||||
dt = datetime.fromisoformat(cleaned)
|
||||
# Ensure timezone-aware (UTC)
|
||||
if dt.tzinfo is None:
|
||||
dt = dt.replace(tzinfo=timezone.utc)
|
||||
return dt
|
||||
except ValueError:
|
||||
return None
|
||||
@@ -0,0 +1,291 @@
|
||||
"""Airflow REST API client with v1/experimental auto-detection."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from urllib.parse import quote
|
||||
|
||||
import requests
|
||||
|
||||
from .config import AirflowConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AirflowAPIError(Exception):
|
||||
"""Raised when Airflow API returns an unexpected error."""
|
||||
|
||||
|
||||
class AirflowClient:
|
||||
"""Client for Apache Airflow REST API.
|
||||
|
||||
Supports both stable API v1 (/api/v1/) and experimental API
|
||||
(/api/experimental/). Auto-detects which is available at startup.
|
||||
"""
|
||||
|
||||
def __init__(self, config: AirflowConfig):
|
||||
self._config = config
|
||||
self._session = requests.Session()
|
||||
self._session.auth = (config.username, config.password)
|
||||
self._session.verify = config.verify_ssl
|
||||
self._session.headers.update({"Content-Type": "application/json"})
|
||||
self._base_url = config.base_url.rstrip("/")
|
||||
self._api_prefix: str | None = None
|
||||
self._api_version: str | None = None
|
||||
self._last_request_time: float = 0
|
||||
logger.debug(
|
||||
"AirflowClient initialized: url=%s, user=%s, timeout=%d, ssl=%s",
|
||||
self._base_url, config.username, config.timeout, config.verify_ssl,
|
||||
)
|
||||
|
||||
def detect_api_version(self):
|
||||
"""Detect available API version.
|
||||
|
||||
Tries v1 stable API first, falls back to experimental.
|
||||
Raises AirflowAPIError if neither is available.
|
||||
"""
|
||||
if self._config.api_version != "auto":
|
||||
if self._config.api_version == "v1":
|
||||
self._api_prefix = f"{self._base_url}/api/v1"
|
||||
self._api_version = "v1"
|
||||
else:
|
||||
self._api_prefix = f"{self._base_url}/api/experimental"
|
||||
self._api_version = "experimental"
|
||||
logger.info("Using configured API version: %s", self._api_version)
|
||||
return
|
||||
|
||||
# Try v1 stable API
|
||||
url = f"{self._base_url}/api/v1/health"
|
||||
logger.debug("Probing API v1: GET %s", url)
|
||||
try:
|
||||
resp = self._session.get(url, timeout=self._config.timeout)
|
||||
logger.debug("Probe v1 response: status=%d, body=%s", resp.status_code, resp.text[:200])
|
||||
if resp.status_code == 200:
|
||||
self._api_prefix = f"{self._base_url}/api/v1"
|
||||
self._api_version = "v1"
|
||||
logger.info("Detected Airflow API v1 (stable)")
|
||||
return
|
||||
except requests.RequestException as e:
|
||||
logger.debug("Probe v1 failed: %s", e)
|
||||
|
||||
# Try experimental API
|
||||
url = f"{self._base_url}/api/experimental/test"
|
||||
logger.debug("Probing experimental API: GET %s", url)
|
||||
try:
|
||||
resp = self._session.get(url, timeout=self._config.timeout)
|
||||
logger.debug("Probe experimental response: status=%d", resp.status_code)
|
||||
if resp.status_code == 200:
|
||||
self._api_prefix = f"{self._base_url}/api/experimental"
|
||||
self._api_version = "experimental"
|
||||
logger.info("Detected Airflow experimental API")
|
||||
return
|
||||
except requests.RequestException as e:
|
||||
logger.debug("Probe experimental failed: %s", e)
|
||||
|
||||
raise AirflowAPIError(
|
||||
f"Cannot connect to Airflow API at {self._base_url}. "
|
||||
"Tried /api/v1/health and /api/experimental/test"
|
||||
)
|
||||
|
||||
def _rate_limit(self):
|
||||
"""Enforce minimum delay between API requests."""
|
||||
elapsed = time.time() - self._last_request_time
|
||||
if elapsed < self._config.request_delay:
|
||||
wait = self._config.request_delay - elapsed
|
||||
logger.debug("Rate limiting: waiting %.2fs", wait)
|
||||
time.sleep(wait)
|
||||
|
||||
def _request(self, method: str, path: str, **kwargs) -> requests.Response:
|
||||
"""Make an API request with error handling and rate limiting."""
|
||||
if not self._api_prefix:
|
||||
raise AirflowAPIError("API version not detected. Call detect_api_version() first")
|
||||
|
||||
self._rate_limit()
|
||||
url = f"{self._api_prefix}{path}"
|
||||
|
||||
# Log request details
|
||||
params = kwargs.get("params")
|
||||
body = kwargs.get("json")
|
||||
logger.debug(
|
||||
"API request: %s %s params=%s body=%s",
|
||||
method, url, params, json.dumps(body) if body else None,
|
||||
)
|
||||
|
||||
t0 = time.monotonic()
|
||||
try:
|
||||
resp = self._session.request(
|
||||
method, url, timeout=self._config.timeout, **kwargs
|
||||
)
|
||||
self._last_request_time = time.time()
|
||||
elapsed_ms = (time.monotonic() - t0) * 1000
|
||||
|
||||
# Log response details
|
||||
logger.debug(
|
||||
"API response: %s %s → %d (%dms) body=%s",
|
||||
method, path, resp.status_code, elapsed_ms,
|
||||
resp.text[:500] if resp.text else "<empty>",
|
||||
)
|
||||
|
||||
resp.raise_for_status()
|
||||
return resp
|
||||
except requests.ConnectionError as e:
|
||||
raise AirflowAPIError(f"Connection error to {url}: {e}") from e
|
||||
except requests.Timeout as e:
|
||||
elapsed_ms = (time.monotonic() - t0) * 1000
|
||||
raise AirflowAPIError(f"Timeout after {elapsed_ms:.0f}ms calling {url}: {e}") from e
|
||||
except requests.HTTPError as e:
|
||||
status = resp.status_code
|
||||
body = resp.text[:500]
|
||||
raise AirflowAPIError(f"HTTP {status} from {url}: {body}") from e
|
||||
|
||||
def get_enabled_dags(self) -> list[dict]:
|
||||
"""Get list of enabled (active and not paused) DAGs.
|
||||
|
||||
Returns list of dicts with at least 'dag_id' key.
|
||||
"""
|
||||
logger.debug("Fetching enabled DAGs (api_version=%s)", self._api_version)
|
||||
if self._api_version == "v1":
|
||||
return self._get_enabled_dags_v1()
|
||||
return self._get_enabled_dags_experimental()
|
||||
|
||||
def _get_enabled_dags_v1(self) -> list[dict]:
|
||||
"""Fetch DAGs via stable v1 API with pagination."""
|
||||
dags = []
|
||||
offset = 0
|
||||
limit = 100
|
||||
|
||||
while True:
|
||||
logger.debug("Fetching DAGs page: offset=%d, limit=%d", offset, limit)
|
||||
resp = self._request(
|
||||
"GET", "/dags",
|
||||
params={"limit": limit, "offset": offset, "only_active": True},
|
||||
)
|
||||
data = resp.json()
|
||||
page_dags = data.get("dags", [])
|
||||
total = data.get("total_entries", 0)
|
||||
|
||||
active_count = 0
|
||||
paused_count = 0
|
||||
for dag in page_dags:
|
||||
if not dag.get("is_paused", True):
|
||||
dags.append(dag)
|
||||
active_count += 1
|
||||
else:
|
||||
paused_count += 1
|
||||
|
||||
logger.debug(
|
||||
"DAGs page result: total_entries=%d, page_size=%d, "
|
||||
"active=%d, paused=%d",
|
||||
total, len(page_dags), active_count, paused_count,
|
||||
)
|
||||
|
||||
offset += limit
|
||||
if offset >= total or not page_dags:
|
||||
break
|
||||
|
||||
logger.debug(
|
||||
"All enabled DAGs: %s",
|
||||
[d["dag_id"] for d in dags],
|
||||
)
|
||||
return dags
|
||||
|
||||
def _get_enabled_dags_experimental(self) -> list[dict]:
|
||||
"""Fetch DAGs via experimental API (limited info)."""
|
||||
resp = self._request("GET", "/dags")
|
||||
return resp.json() if resp.status_code == 200 else []
|
||||
|
||||
def get_dag_runs(self, dag_id: str, states: list[str] | None = None) -> list[dict]:
|
||||
"""Get recent DAG runs, optionally filtered by state."""
|
||||
encoded_dag_id = quote(dag_id, safe="")
|
||||
logger.debug("Fetching DAG runs: dag_id=%s, states=%s", dag_id, states)
|
||||
|
||||
if self._api_version == "v1":
|
||||
if states:
|
||||
resp = self._request(
|
||||
"GET",
|
||||
f"/dags/{encoded_dag_id}/dagRuns",
|
||||
params=[("limit", 25), ("order_by", "-start_date")]
|
||||
+ [("state", s) for s in states],
|
||||
)
|
||||
else:
|
||||
resp = self._request(
|
||||
"GET",
|
||||
f"/dags/{encoded_dag_id}/dagRuns",
|
||||
params={"order_by": "-start_date", "limit": 25},
|
||||
)
|
||||
runs = resp.json().get("dag_runs", [])
|
||||
else:
|
||||
resp = self._request("GET", f"/dags/{dag_id}/dag_runs")
|
||||
runs = resp.json() if isinstance(resp.json(), list) else []
|
||||
if states:
|
||||
runs = [r for r in runs if r.get("state") in states]
|
||||
|
||||
logger.debug(
|
||||
"DAG runs for %s: count=%d, runs=[%s]",
|
||||
dag_id, len(runs),
|
||||
", ".join(
|
||||
f"{r.get('dag_run_id', '?')}({r.get('state', '?')})"
|
||||
for r in runs
|
||||
),
|
||||
)
|
||||
return runs
|
||||
|
||||
def get_task_instances(self, dag_id: str, dag_run_id: str) -> list[dict]:
|
||||
"""Get task instances for a specific DAG run."""
|
||||
encoded_dag_id = quote(dag_id, safe="")
|
||||
encoded_run_id = quote(dag_run_id, safe="")
|
||||
|
||||
logger.debug("Fetching task instances: %s/%s", dag_id, dag_run_id)
|
||||
|
||||
if self._api_version == "v1":
|
||||
resp = self._request(
|
||||
"GET",
|
||||
f"/dags/{encoded_dag_id}/dagRuns/{encoded_run_id}/taskInstances",
|
||||
)
|
||||
tasks = resp.json().get("task_instances", [])
|
||||
logger.debug(
|
||||
"Task instances for %s/%s: count=%d, states=[%s]",
|
||||
dag_id, dag_run_id, len(tasks),
|
||||
", ".join(
|
||||
f"{t.get('task_id', '?')}({t.get('state', '?')})"
|
||||
for t in tasks
|
||||
),
|
||||
)
|
||||
return tasks
|
||||
else:
|
||||
logger.warning(
|
||||
"Task instances not supported in experimental API for %s/%s",
|
||||
dag_id, dag_run_id,
|
||||
)
|
||||
return []
|
||||
|
||||
def clear_dag_run(self, dag_id: str, dag_run_id: str) -> bool:
|
||||
"""Clear (retry) a DAG run by resetting failed task instances.
|
||||
|
||||
Returns True on success, False on failure.
|
||||
"""
|
||||
encoded_dag_id = quote(dag_id, safe="")
|
||||
encoded_run_id = quote(dag_run_id, safe="")
|
||||
|
||||
if self._api_version != "v1":
|
||||
logger.warning(
|
||||
"Clearing DAG runs not supported in experimental API for %s/%s",
|
||||
dag_id, dag_run_id,
|
||||
)
|
||||
return False
|
||||
|
||||
logger.debug(
|
||||
"Clearing DAG run: %s/%s (only_failed=True)",
|
||||
dag_id, dag_run_id,
|
||||
)
|
||||
try:
|
||||
self._request(
|
||||
"POST",
|
||||
f"/dags/{encoded_dag_id}/dagRuns/{encoded_run_id}/clear",
|
||||
json={"dry_run": False, "only_failed": True},
|
||||
)
|
||||
logger.info("Cleared DAG run %s/%s for retry", dag_id, dag_run_id)
|
||||
return True
|
||||
except AirflowAPIError as e:
|
||||
logger.error("Failed to clear DAG run %s/%s: %s", dag_id, dag_run_id, e)
|
||||
return False
|
||||
@@ -0,0 +1,100 @@
|
||||
"""Configuration loader for Airflow DAG Monitor."""
|
||||
|
||||
import dataclasses
|
||||
import pathlib
|
||||
|
||||
import yaml
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class AirflowConfig:
|
||||
"""Airflow API connection settings."""
|
||||
|
||||
base_url: str = "http://localhost:8080"
|
||||
username: str = "airflow"
|
||||
password: str = "airflow"
|
||||
api_version: str = "auto" # "auto", "v1", or "experimental"
|
||||
timeout: int = 30
|
||||
verify_ssl: bool = True
|
||||
request_delay: float = 0.5 # delay between API calls (seconds)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class MonitorConfig:
|
||||
"""Monitoring behavior settings."""
|
||||
|
||||
cycle_interval: int = 300 # seconds between full cycles
|
||||
long_running_threshold: int = 1800 # 30 minutes
|
||||
retry_wait: int = 120 # wait after retry before re-check
|
||||
max_retries: int = 1
|
||||
state_file: str = "/var/lib/airflow-monitor/state.json"
|
||||
lock_file: str = "/var/run/airflow-monitor.lock"
|
||||
state_max_age: int = 86400 # 24 hours
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class ZabbixConfig:
|
||||
"""Zabbix agent integration settings.
|
||||
|
||||
Monitor writes data to files, Zabbix agent reads via UserParameter.
|
||||
"""
|
||||
|
||||
enabled: bool = True
|
||||
data_dir: str = "/var/lib/airflow-monitor"
|
||||
problems_file: str = "problems.json"
|
||||
heartbeat_file: str = "heartbeat"
|
||||
status_file: str = "status.json"
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class LoggingConfig:
|
||||
"""Logging settings."""
|
||||
|
||||
level: str = "INFO"
|
||||
file: str = "/var/log/airflow-monitor/monitor.log"
|
||||
max_bytes: int = 10_485_760 # 10 MB
|
||||
backup_count: int = 5
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class AppConfig:
|
||||
"""Top-level application configuration."""
|
||||
|
||||
airflow: AirflowConfig
|
||||
monitor: MonitorConfig
|
||||
zabbix: ZabbixConfig
|
||||
logging: LoggingConfig
|
||||
|
||||
|
||||
def _build_dataclass(cls, data: dict):
|
||||
"""Build a dataclass instance from a dict, ignoring unknown keys."""
|
||||
if data is None:
|
||||
return cls()
|
||||
fields = {f.name for f in dataclasses.fields(cls)}
|
||||
filtered = {k: v for k, v in data.items() if k in fields}
|
||||
return cls(**filtered)
|
||||
|
||||
|
||||
def load_config(path: str) -> AppConfig:
|
||||
"""Load configuration from a YAML file.
|
||||
|
||||
Missing sections fall back to defaults.
|
||||
Raises FileNotFoundError if the file does not exist.
|
||||
Raises ValueError on invalid YAML.
|
||||
"""
|
||||
config_path = pathlib.Path(path)
|
||||
if not config_path.exists():
|
||||
raise FileNotFoundError(f"Config file not found: {path}")
|
||||
|
||||
with open(config_path, "r", encoding="utf-8") as f:
|
||||
raw = yaml.safe_load(f)
|
||||
|
||||
if not isinstance(raw, dict):
|
||||
raise ValueError(f"Config file must be a YAML mapping, got {type(raw).__name__}")
|
||||
|
||||
return AppConfig(
|
||||
airflow=_build_dataclass(AirflowConfig, raw.get("airflow")),
|
||||
monitor=_build_dataclass(MonitorConfig, raw.get("monitor")),
|
||||
zabbix=_build_dataclass(ZabbixConfig, raw.get("zabbix")),
|
||||
logging=_build_dataclass(LoggingConfig, raw.get("logging")),
|
||||
)
|
||||
@@ -0,0 +1,333 @@
|
||||
"""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)
|
||||
@@ -0,0 +1,250 @@
|
||||
"""Main monitoring loop - orchestrates the full monitoring cycle."""
|
||||
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
|
||||
from .actions import ActionHandler, DataExporter
|
||||
from .analyzer import DagAnalyzer
|
||||
from .client import AirflowAPIError, AirflowClient
|
||||
from .config import AppConfig
|
||||
from .state import StateManager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Monitor:
|
||||
"""Orchestrates the Airflow DAG monitoring cycle.
|
||||
|
||||
Cycle flow:
|
||||
1. Load state
|
||||
2. Fetch enabled DAGs
|
||||
3. Get active runs for each DAG
|
||||
4. Analyze for issues (failed, long-running)
|
||||
5. Retry where possible
|
||||
6. Wait and re-check retried runs
|
||||
7. Alert via Zabbix for unresolved issues
|
||||
8. Send heartbeat, cleanup, save state
|
||||
"""
|
||||
|
||||
def __init__(self, config: AppConfig, shutdown_event: threading.Event):
|
||||
self._config = config
|
||||
self._shutdown = shutdown_event
|
||||
self._client = AirflowClient(config.airflow)
|
||||
self._state = StateManager(config.monitor.state_file, config.monitor.state_max_age)
|
||||
self._analyzer = DagAnalyzer(config.monitor.long_running_threshold)
|
||||
self._exporter = DataExporter(config.zabbix)
|
||||
self._actions = ActionHandler(
|
||||
self._client, self._state, self._exporter, config.monitor,
|
||||
)
|
||||
self._cycle_count = 0
|
||||
|
||||
def run(self):
|
||||
"""Main loop. Runs until shutdown_event is set."""
|
||||
logger.info(
|
||||
"Starting Airflow DAG monitor (cycle=%ds, threshold=%ds, retries=%d, retry_wait=%ds)",
|
||||
self._config.monitor.cycle_interval,
|
||||
self._config.monitor.long_running_threshold,
|
||||
self._config.monitor.max_retries,
|
||||
self._config.monitor.retry_wait,
|
||||
)
|
||||
|
||||
try:
|
||||
self._client.detect_api_version()
|
||||
except AirflowAPIError as e:
|
||||
logger.critical("Cannot connect to Airflow API: %s", e)
|
||||
return
|
||||
|
||||
while not self._shutdown.is_set():
|
||||
try:
|
||||
self._cycle()
|
||||
except Exception:
|
||||
logger.exception("Unhandled error in monitoring cycle")
|
||||
|
||||
# Interruptible sleep between cycles
|
||||
logger.debug(
|
||||
"Sleeping %ds until next cycle",
|
||||
self._config.monitor.cycle_interval,
|
||||
)
|
||||
self._shutdown.wait(timeout=self._config.monitor.cycle_interval)
|
||||
|
||||
logger.info("Shutting down gracefully")
|
||||
self._state.save()
|
||||
|
||||
def _cycle(self):
|
||||
"""Execute one complete monitoring cycle."""
|
||||
self._cycle_count += 1
|
||||
cycle_start = time.monotonic()
|
||||
logger.info("=== Monitoring cycle #%d started ===", self._cycle_count)
|
||||
self._state.load()
|
||||
|
||||
# --- Phase 1: Collect issues ---
|
||||
phase_start = time.monotonic()
|
||||
all_issues = []
|
||||
try:
|
||||
dags = self._client.get_enabled_dags()
|
||||
except AirflowAPIError as e:
|
||||
logger.error("Failed to fetch DAG list: %s", e)
|
||||
self._exporter.export_heartbeat()
|
||||
return
|
||||
|
||||
logger.info(
|
||||
"Phase 1/5 [Collect]: found %d enabled DAGs (%.1fs)",
|
||||
len(dags), time.monotonic() - phase_start,
|
||||
)
|
||||
|
||||
for i, dag in enumerate(dags, 1):
|
||||
if self._shutdown.is_set():
|
||||
return
|
||||
|
||||
dag_id = dag.get("dag_id", "")
|
||||
if not dag_id:
|
||||
continue
|
||||
|
||||
logger.debug("Processing DAG %d/%d: %s", i, len(dags), dag_id)
|
||||
|
||||
try:
|
||||
runs = self._client.get_dag_runs(dag_id, states=["running", "failed"])
|
||||
except AirflowAPIError as e:
|
||||
logger.warning("Failed to fetch runs for DAG '%s': %s", dag_id, e)
|
||||
continue
|
||||
|
||||
if not runs:
|
||||
logger.debug(" No active/failed runs for %s", dag_id)
|
||||
continue
|
||||
|
||||
logger.debug(" Found %d active/failed runs for %s", len(runs), dag_id)
|
||||
|
||||
issues = self._analyzer.analyze_dag_runs(
|
||||
dag_id, runs,
|
||||
task_instances_fn=self._client.get_task_instances,
|
||||
)
|
||||
all_issues.extend(issues)
|
||||
|
||||
if all_issues:
|
||||
logger.warning(
|
||||
"Phase 1 result: %d issues found: %s",
|
||||
len(all_issues),
|
||||
", ".join(
|
||||
f"{i.dag_id}/{i.dag_run_id}({i.issue_type}, {i.duration_seconds:.0f}s)"
|
||||
for i in all_issues
|
||||
),
|
||||
)
|
||||
else:
|
||||
logger.info("Phase 1 result: no issues found across %d DAGs", len(dags))
|
||||
|
||||
# --- Phase 2: Handle issues (retry or mark for alert) ---
|
||||
phase_start = time.monotonic()
|
||||
needs_recheck = []
|
||||
for issue in all_issues:
|
||||
action = self._actions.handle_issue(issue)
|
||||
logger.info(
|
||||
"Phase 2/5 [Handle]: DAG %s run %s → %s (type=%s, duration=%.0fs)",
|
||||
issue.dag_id, issue.dag_run_id, action,
|
||||
issue.issue_type, issue.duration_seconds,
|
||||
)
|
||||
if action == "retried":
|
||||
needs_recheck.append(issue)
|
||||
|
||||
logger.info(
|
||||
"Phase 2/5 [Handle]: processed %d issues, %d retried, (%.1fs)",
|
||||
len(all_issues), len(needs_recheck),
|
||||
time.monotonic() - phase_start,
|
||||
)
|
||||
|
||||
# --- Phase 3: Wait and re-check retried runs ---
|
||||
if needs_recheck and not self._shutdown.is_set():
|
||||
wait_time = self._config.monitor.retry_wait
|
||||
logger.info(
|
||||
"Phase 3/5 [Wait]: waiting %ds to re-check %d retried runs",
|
||||
wait_time, len(needs_recheck),
|
||||
)
|
||||
self._shutdown.wait(timeout=wait_time)
|
||||
|
||||
if not self._shutdown.is_set():
|
||||
phase_start = time.monotonic()
|
||||
still_failing = self._recheck_retried(needs_recheck)
|
||||
# Replace all_issues with only the still-failing ones for alerting
|
||||
# Keep non-retried issues that need alerting
|
||||
non_retried_issues = [
|
||||
i for i in all_issues if i not in needs_recheck
|
||||
]
|
||||
all_issues = non_retried_issues + still_failing
|
||||
logger.info(
|
||||
"Phase 3/5 [Recheck]: %d/%d still failing after retry (%.1fs)",
|
||||
len(still_failing), len(needs_recheck),
|
||||
time.monotonic() - phase_start,
|
||||
)
|
||||
else:
|
||||
logger.debug("Phase 3/5 [Wait]: skipped (no retried runs)")
|
||||
|
||||
# --- Phase 4: Alert via Zabbix ---
|
||||
phase_start = time.monotonic()
|
||||
self._actions.collect_and_alert(all_issues)
|
||||
logger.info(
|
||||
"Phase 4/5 [Alert]: alert phase complete (%.1fs)",
|
||||
time.monotonic() - phase_start,
|
||||
)
|
||||
|
||||
# --- Phase 5: Housekeeping ---
|
||||
phase_start = time.monotonic()
|
||||
self._exporter.export_heartbeat()
|
||||
self._state.purge_old_entries()
|
||||
self._state.save()
|
||||
|
||||
cycle_elapsed = time.monotonic() - cycle_start
|
||||
self._exporter.export_status(
|
||||
self._cycle_count, len(dags), len(all_issues), cycle_elapsed,
|
||||
)
|
||||
logger.info(
|
||||
"=== Monitoring cycle #%d complete: %d DAGs checked, "
|
||||
"%d issues, cycle_time=%.1fs ===",
|
||||
self._cycle_count, len(dags), len(all_issues), cycle_elapsed,
|
||||
)
|
||||
|
||||
def _recheck_retried(self, retried_issues: list) -> list:
|
||||
"""Re-check DAG runs that were retried.
|
||||
|
||||
Returns list of DagIssue objects that are still failing.
|
||||
"""
|
||||
still_failing = []
|
||||
|
||||
for issue in retried_issues:
|
||||
if self._shutdown.is_set():
|
||||
break
|
||||
|
||||
logger.debug("Re-checking retried run: %s/%s", issue.dag_id, issue.dag_run_id)
|
||||
|
||||
try:
|
||||
runs = self._client.get_dag_runs(
|
||||
issue.dag_id, states=["running", "failed"],
|
||||
)
|
||||
except AirflowAPIError as e:
|
||||
logger.warning(
|
||||
"Failed to re-check DAG '%s': %s. Treating as still failing.",
|
||||
issue.dag_id, e,
|
||||
)
|
||||
still_failing.append(issue)
|
||||
continue
|
||||
|
||||
recheck_issues = self._analyzer.analyze_dag_runs(
|
||||
issue.dag_id, runs,
|
||||
task_instances_fn=self._client.get_task_instances,
|
||||
)
|
||||
|
||||
# Check if the same dag_run_id still has problems
|
||||
for ri in recheck_issues:
|
||||
if ri.dag_run_id == issue.dag_run_id:
|
||||
logger.warning(
|
||||
"DAG %s/%s still failing after retry: type=%s, duration=%.0fs",
|
||||
ri.dag_id, ri.dag_run_id, ri.issue_type, ri.duration_seconds,
|
||||
)
|
||||
still_failing.append(ri)
|
||||
break
|
||||
else:
|
||||
logger.info(
|
||||
"DAG %s/%s recovered after retry",
|
||||
issue.dag_id, issue.dag_run_id,
|
||||
)
|
||||
|
||||
return still_failing
|
||||
@@ -0,0 +1,164 @@
|
||||
"""Persistent state manager for tracking DAG run retries and alerts."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import pathlib
|
||||
import tempfile
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class StateManager:
|
||||
"""Manages persistent state in a JSON file.
|
||||
|
||||
State tracks retry counts and alert flags per DAG run to ensure
|
||||
idempotent behavior across monitoring cycles.
|
||||
"""
|
||||
|
||||
def __init__(self, state_file: str, max_age: int = 86400):
|
||||
self._path = pathlib.Path(state_file)
|
||||
self._max_age = max_age
|
||||
self._data: dict = {"version": 1, "entries": {}}
|
||||
logger.debug(
|
||||
"StateManager initialized: file=%s, max_age=%ds",
|
||||
state_file, max_age,
|
||||
)
|
||||
|
||||
def load(self):
|
||||
"""Load state from disk. Start fresh if missing or corrupt."""
|
||||
if not self._path.exists():
|
||||
logger.debug("State file not found at %s, starting with empty state", self._path)
|
||||
self._data = {"version": 1, "entries": {}}
|
||||
return
|
||||
|
||||
try:
|
||||
with open(self._path, "r", encoding="utf-8") as f:
|
||||
raw = f.read()
|
||||
self._data = json.loads(raw)
|
||||
if not isinstance(self._data.get("entries"), dict):
|
||||
raise ValueError("Invalid state structure: 'entries' is not a dict")
|
||||
entry_count = len(self._data["entries"])
|
||||
logger.debug(
|
||||
"State loaded: %d entries, file_size=%d bytes",
|
||||
entry_count, len(raw),
|
||||
)
|
||||
if entry_count > 0:
|
||||
logger.debug(
|
||||
"State entries: %s",
|
||||
", ".join(
|
||||
f"{k}(retries={v.get('retry_count', 0)}, alerted={v.get('alerted', False)})"
|
||||
for k, v in self._data["entries"].items()
|
||||
),
|
||||
)
|
||||
except (json.JSONDecodeError, ValueError, KeyError) as e:
|
||||
logger.warning("Corrupt state file %s, starting fresh: %s", self._path, e)
|
||||
self._data = {"version": 1, "entries": {}}
|
||||
|
||||
def save(self):
|
||||
"""Save state to disk atomically (write tmp -> rename)."""
|
||||
self._path.parent.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
content = json.dumps(self._data, indent=2, ensure_ascii=False)
|
||||
fd, tmp_path = tempfile.mkstemp(
|
||||
dir=str(self._path.parent), suffix=".tmp"
|
||||
)
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
||||
f.write(content)
|
||||
os.replace(tmp_path, str(self._path))
|
||||
logger.debug(
|
||||
"State saved: %d entries, %d bytes → %s",
|
||||
len(self._data["entries"]), len(content), self._path,
|
||||
)
|
||||
except OSError as e:
|
||||
logger.error("Failed to save state to %s: %s", self._path, e)
|
||||
# Clean up temp file if it exists
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _key(dag_id: str, dag_run_id: str) -> str:
|
||||
return f"{dag_id}::{dag_run_id}"
|
||||
|
||||
@staticmethod
|
||||
def _now_iso() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
def get_entry(self, dag_id: str, dag_run_id: str) -> dict | None:
|
||||
"""Get stored entry for a DAG run, or None."""
|
||||
return self._data["entries"].get(self._key(dag_id, dag_run_id))
|
||||
|
||||
def ensure_entry(self, dag_id: str, dag_run_id: str) -> dict:
|
||||
"""Get or create an entry for a DAG run."""
|
||||
key = self._key(dag_id, dag_run_id)
|
||||
if key not in self._data["entries"]:
|
||||
self._data["entries"][key] = {
|
||||
"dag_id": dag_id,
|
||||
"dag_run_id": dag_run_id,
|
||||
"retry_count": 0,
|
||||
"last_retry_time": None,
|
||||
"alerted": False,
|
||||
"first_seen": self._now_iso(),
|
||||
"last_seen": self._now_iso(),
|
||||
}
|
||||
logger.debug("State: new entry created for %s", key)
|
||||
else:
|
||||
self._data["entries"][key]["last_seen"] = self._now_iso()
|
||||
logger.debug("State: updated last_seen for %s", key)
|
||||
return self._data["entries"][key]
|
||||
|
||||
def get_retry_count(self, dag_id: str, dag_run_id: str) -> int:
|
||||
"""Get current retry count for a DAG run."""
|
||||
entry = self.get_entry(dag_id, dag_run_id)
|
||||
count = entry["retry_count"] if entry else 0
|
||||
logger.debug("State: retry_count for %s/%s = %d", dag_id, dag_run_id, count)
|
||||
return count
|
||||
|
||||
def increment_retry(self, dag_id: str, dag_run_id: str):
|
||||
"""Increment retry counter and record time."""
|
||||
entry = self.ensure_entry(dag_id, dag_run_id)
|
||||
old_count = entry["retry_count"]
|
||||
entry["retry_count"] += 1
|
||||
entry["last_retry_time"] = self._now_iso()
|
||||
logger.debug(
|
||||
"State: retry_count for %s/%s incremented %d → %d",
|
||||
dag_id, dag_run_id, old_count, entry["retry_count"],
|
||||
)
|
||||
|
||||
def mark_alerted(self, dag_id: str, dag_run_id: str):
|
||||
"""Mark a DAG run as alerted (avoid duplicate alerts)."""
|
||||
entry = self.ensure_entry(dag_id, dag_run_id)
|
||||
entry["alerted"] = True
|
||||
logger.debug("State: marked alerted for %s/%s", dag_id, dag_run_id)
|
||||
|
||||
def is_alerted(self, dag_id: str, dag_run_id: str) -> bool:
|
||||
"""Check if alert was already sent for this DAG run."""
|
||||
entry = self.get_entry(dag_id, dag_run_id)
|
||||
return entry["alerted"] if entry else False
|
||||
|
||||
def purge_old_entries(self):
|
||||
"""Remove entries older than max_age seconds."""
|
||||
cutoff = time.time() - self._max_age
|
||||
keys_to_remove = []
|
||||
|
||||
for key, entry in self._data["entries"].items():
|
||||
last_seen = entry.get("last_seen", entry.get("first_seen", ""))
|
||||
try:
|
||||
ts = datetime.fromisoformat(last_seen).timestamp()
|
||||
if ts < cutoff:
|
||||
keys_to_remove.append(key)
|
||||
except (ValueError, TypeError):
|
||||
keys_to_remove.append(key)
|
||||
|
||||
for key in keys_to_remove:
|
||||
logger.debug("State: purging old entry %s", key)
|
||||
del self._data["entries"][key]
|
||||
|
||||
if keys_to_remove:
|
||||
logger.info("Purged %d old state entries (max_age=%ds)", len(keys_to_remove), self._max_age)
|
||||
else:
|
||||
logger.debug("State: no entries to purge")
|
||||
Reference in New Issue
Block a user