Set up Apache Airflow performance monitoring with DataDog agent integration and custom dashboards

#!/usr/bin/env bash
set -euo pipefail

# Apache Airflow DataDog Monitoring Setup Script
# Configures DataDog agent integration with Apache Airflow for comprehensive monitoring

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Default values
DD_API_KEY="${1:-}"
AIRFLOW_USER="${2:-airflow}"
AIRFLOW_HOME="${3:-/opt/airflow}"
AIRFLOW_WEB_USER="${4:-admin}"
AIRFLOW_WEB_PASS="${5:-admin}"

usage() {
    echo "Usage: $0 <DD_API_KEY> [airflow_user] [airflow_home] [web_user] [web_pass]"
    echo "  DD_API_KEY: DataDog API key (required)"
    echo "  airflow_user: Airflow system user (default: airflow)"
    echo "  airflow_home: Airflow installation directory (default: /opt/airflow)"
    echo "  web_user: Airflow web UI username (default: admin)"
    echo "  web_pass: Airflow web UI password (default: admin)"
    exit 1
}

log() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

error() {
    echo -e "${RED}[ERROR]${NC} $1" >&2
}

cleanup() {
    if [ $? -ne 0 ]; then
        error "Installation failed. Cleaning up..."
        systemctl stop datadog-agent 2>/dev/null || true
        rm -f /etc/datadog-agent/conf.d/airflow.d/conf.yaml
        rm -f /opt/datadog-agent/embedded/bin/airflow_metrics.py
    fi
}

trap cleanup ERR

# Validate arguments
if [ -z "$DD_API_KEY" ]; then
    usage
fi

# Check if running as root
if [ "$EUID" -ne 0 ]; then
    error "This script must be run as root"
    exit 1
fi

# Detect distribution
if [ -f /etc/os-release ]; then
    . /etc/os-release
    case "$ID" in
        ubuntu|debian)
            PKG_MGR="apt"
            PKG_INSTALL="apt install -y"
            PKG_UPDATE="apt update"
            ;;
        almalinux|rocky|centos|rhel|ol)
            PKG_MGR="dnf"
            PKG_INSTALL="dnf install -y"
            PKG_UPDATE="dnf check-update || true"
            ;;
        fedora)
            PKG_MGR="dnf"
            PKG_INSTALL="dnf install -y"
            PKG_UPDATE="dnf check-update || true"
            ;;
        amzn)
            PKG_MGR="yum"
            PKG_INSTALL="yum install -y"
            PKG_UPDATE="yum check-update || true"
            ;;
        *)
            error "Unsupported distribution: $ID"
            exit 1
            ;;
    esac
else
    error "Cannot detect distribution - /etc/os-release not found"
    exit 1
fi

log "Detected distribution: $ID using $PKG_MGR"

echo "[1/8] Updating package repositories..."
$PKG_UPDATE

echo "[2/8] Installing prerequisites..."
$PKG_INSTALL curl wget python3 python3-pip

echo "[3/8] Installing DataDog agent..."
DD_API_KEY="$DD_API_KEY" bash -c "$(curl -L https://s3.amazonaws.com/dd-agent/scripts/install_script.sh)"

systemctl enable datadog-agent
systemctl start datadog-agent

echo "[4/8] Configuring DataDog agent for StatsD..."
cat > /etc/datadog-agent/datadog.yaml.tmp << EOF
api_key: $DD_API_KEY
site: datadoghq.com

# DogStatsD configuration
dogstatsd_config:
  enabled: true
  bind_host: localhost
  port: 8125
  non_local_traffic: false

# Tags for all metrics
tags:
  - datacenter:$(curl -s http://169.254.169.254/latest/meta-data/placement/availability-zone 2>/dev/null || echo "unknown")
  - environment:production
  - service:airflow

logs_enabled: true
process_config:
  enabled: true
EOF

# Preserve existing API key and site if datadog.yaml exists
if [ -f /etc/datadog-agent/datadog.yaml ]; then
    cp /etc/datadog-agent/datadog.yaml /etc/datadog-agent/datadog.yaml.bak
fi

mv /etc/datadog-agent/datadog.yaml.tmp /etc/datadog-agent/datadog.yaml
chown dd-agent:dd-agent /etc/datadog-agent/datadog.yaml
chmod 640 /etc/datadog-agent/datadog.yaml

echo "[5/8] Creating Airflow integration configuration..."
mkdir -p /etc/datadog-agent/conf.d/airflow.d
cat > /etc/datadog-agent/conf.d/airflow.d/conf.yaml << EOF
init_config:

instances:
  - url: http://localhost:8080
    username: $AIRFLOW_WEB_USER
    password: $AIRFLOW_WEB_PASS
    tags:
      - environment:production
      - airflow_cluster:main
    collect_health_metrics: true
    collect_task_metrics: true
    collect_dag_metrics: true
EOF

chown dd-agent:dd-agent /etc/datadog-agent/conf.d/airflow.d/conf.yaml
chmod 644 /etc/datadog-agent/conf.d/airflow.d/conf.yaml

echo "[6/8] Configuring Airflow metrics collection..."
if [ -f "$AIRFLOW_HOME/airflow.cfg" ]; then
    # Backup existing config
    cp "$AIRFLOW_HOME/airflow.cfg" "$AIRFLOW_HOME/airflow.cfg.bak.$(date +%Y%m%d_%H%M%S)"
    
    # Update metrics configuration
    if grep -q "\[metrics\]" "$AIRFLOW_HOME/airflow.cfg"; then
        # Update existing metrics section
        sed -i '/^\[metrics\]/,/^\[/ {
            s/^statsd_on.*/statsd_on = True/
            s/^statsd_host.*/statsd_host = localhost/
            s/^statsd_port.*/statsd_port = 8125/
            s/^statsd_prefix.*/statsd_prefix = airflow/
            s/^statsd_allow_list.*/statsd_allow_list = scheduler,executor,dagrun,taskinstance/
        }' "$AIRFLOW_HOME/airflow.cfg"
    else
        # Add metrics section
        cat >> "$AIRFLOW_HOME/airflow.cfg" << EOF

[metrics]
statsd_on = True
statsd_host = localhost
statsd_port = 8125
statsd_prefix = airflow
statsd_allow_list = scheduler,executor,dagrun,taskinstance
EOF
    fi
    
    chown $AIRFLOW_USER:$AIRFLOW_USER "$AIRFLOW_HOME/airflow.cfg"
    chmod 644 "$AIRFLOW_HOME/airflow.cfg"
    log "Updated Airflow configuration at $AIRFLOW_HOME/airflow.cfg"
else
    warn "Airflow configuration file not found at $AIRFLOW_HOME/airflow.cfg"
    warn "Please manually add metrics configuration to airflow.cfg"
fi

echo "[7/8] Creating custom metrics collection script..."
mkdir -p /opt/datadog-agent/embedded/bin

cat > /opt/datadog-agent/embedded/bin/airflow_metrics.py << 'EOF'
#!/usr/bin/env python3
import time
import sys
import os
from datetime import datetime, timedelta

try:
    from airflow import settings
    from airflow.models import DagRun, TaskInstance, DagModel
    from sqlalchemy import func
except ImportError:
    print("Airflow not found. Please ensure Airflow is installed and accessible.")
    sys.exit(1)

try:
    import socket
except ImportError:
    print("Socket module not available")
    sys.exit(1)

def send_metric(metric_name, value, metric_type='gauge', host='localhost', port=8125):
    """Send metric to StatsD"""
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        if metric_type == 'gauge':
            message = f"{metric_name}:{value}|g"
        elif metric_type == 'count':
            message = f"{metric_name}:{value}|c"
        sock.sendto(message.encode(), (host, port))
    finally:
        sock.close()

def collect_airflow_metrics():
    try:
        session = settings.Session()
        
        # Count active DAGs
        active_dags = session.query(DagModel).filter(DagModel.is_active == True).count()
        send_metric('airflow.dags.active', active_dags)
        
        # Count running DAG runs
        running_dag_runs = session.query(DagRun).filter(DagRun.state == 'running').count()
        send_metric('airflow.dag_runs.running', running_dag_runs)
        
        # Count failed tasks in last hour
        one_hour_ago = datetime.utcnow() - timedelta(hours=1)
        failed_tasks = session.query(TaskInstance).filter(
            TaskInstance.state == 'failed',
            TaskInstance.end_date >= one_hour_ago
        ).count()
        send_metric('airflow.tasks.failed_last_hour', failed_tasks)
        
        print("Custom metrics collected successfully")
        
    except Exception as e:
        print(f"Error collecting metrics: {e}")
    finally:
        session.close()

if __name__ == "__main__":
    collect_airflow_metrics()
EOF

chmod 755 /opt/datadog-agent/embedded/bin/airflow_metrics.py
chown dd-agent:dd-agent /opt/datadog-agent/embedded/bin/airflow_metrics.py

# Create systemd service for custom metrics
cat > /etc/systemd/system/airflow-datadog-metrics.service << EOF
[Unit]
Description=Airflow DataDog Custom Metrics Collection
After=datadog-agent.service

[Service]
Type=oneshot
User=dd-agent
Environment=AIRFLOW_HOME=$AIRFLOW_HOME
ExecStart=/opt/datadog-agent/embedded/bin/airflow_metrics.py
EOF

# Create timer for regular execution
cat > /etc/systemd/system/airflow-datadog-metrics.timer << EOF
[Unit]
Description=Run Airflow DataDog metrics collection every 5 minutes
Requires=airflow-datadog-metrics.service

[Timer]
OnCalendar=*:0/5
Persistent=true

[Install]
WantedBy=timers.target
EOF

systemctl daemon-reload
systemctl enable airflow-datadog-metrics.timer
systemctl start airflow-datadog-metrics.timer

echo "[8/8] Restarting DataDog agent and verification..."
systemctl restart datadog-agent

# Wait for services to start
sleep 5

# Verify installation
echo "Verifying installation..."
if systemctl is-active --quiet datadog-agent; then
    log "✓ DataDog agent is running"
else
    error "✗ DataDog agent is not running"
    exit 1
fi

if systemctl is-active --quiet airflow-datadog-metrics.timer; then
    log "✓ Custom metrics timer is active"
else
    warn "✗ Custom metrics timer is not active"
fi

if [ -f /etc/datadog-agent/conf.d/airflow.d/conf.yaml ]; then
    log "✓ Airflow integration configured"
else
    error "✗ Airflow integration configuration missing"
fi

log "Installation completed successfully!"
log "DataDog agent is now monitoring your Airflow instance."
log "Check your DataDog dashboard for incoming metrics."
log "Custom metrics will be collected every 5 minutes."