"""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))