"""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.chmod(tmp_path, 0o644) 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, all_issues: list[DagIssue]) -> str: """Process a single issue. Args: issue: The issue to handle. all_issues: All issues in this cycle (to check for running runs of same DAG). Returns: "already_handled" - previously alerted, skip "retried" - retry initiated, needs re-check "needs_alert" - retry exhausted, alert needed "skip_running" - DAG already has a running instance, unsafe to retry """ 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" # SAFETY: Do not retry if this DAG already has a running instance. # Prevents parallel execution of the same DAG (critical for reports). if issue.issue_type == "failed": has_running = any( i.dag_id == dag_id and i.state == "running" for i in all_issues ) if has_running: logger.warning( " → skip_running: DAG %s has a running instance, " "will NOT retry failed run %s to avoid parallel execution", dag_id, run_id, ) return "needs_alert" # 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 all current issues and mark new ones as alerted. Always exports the full list of current problems so Zabbix sees them until they are resolved. Only logs warnings for newly discovered issues. """ # Build export payload for ALL current issues all_problems = [] new_issues = [] for issue in issues: 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"], } all_problems.append(alert_data) if not self._state.is_alerted(issue.dag_id, issue.dag_run_id): new_issues.append(issue) if not all_problems: self._exporter.export_problems([]) logger.info("No issues, exported empty problem list") return # Log new issues at warning level if new_issues: logger.warning( "New problematic DAG runs: %s", ", ".join( f"{i.dag_id}/{i.dag_run_id}({i.issue_type})" for i in new_issues ), ) logger.debug("Alert payload:\n%s", json.dumps(all_problems, indent=2, ensure_ascii=False)) self._exporter.export_problems(all_problems) # Mark new issues as alerted for issue in new_issues: 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 (%d new)", len(all_problems), len(new_issues), )