# In your airflow.cfg or helm values  # Place this configuration in airflow.cfg or as Helm chart values
[metrics]  # Beginning of the metrics configuration section in Airflow
prometheus_enabled = True  # Enable the Prometheus metrics exporter in Airflow
prometheus_endpoint = /metrics  # Set the endpoint path where Prometheus scrapes Airflow metrics

from airflow import DAG  # Import the core DAG class for defining Airflow workflows
from airflow.operators.python import PythonOperator  # Import the operator that runs arbitrary Python functions
from datetime import datetime  # Import datetime for specifying the DAG's start_date parameter
import requests  # Import the HTTP library for scraping metrics from internal endpoints

def scrape_metrics():  # Define the function that will scrape internal endpoints and push to PushGateway
    targets = [  # List of internal service URLs that expose Prometheus metrics
        "http://internal-api:8080/metrics",  # Internal API service metrics endpoint
        "http://batch-worker:9090/metrics",  # Batch worker service metrics endpoint
    ]
    for target in targets:  # Iterate over each target endpoint
        try:  # Wrap scraping in try-except to handle transient network or service failures
            resp = requests.get(target, timeout=10)  # Send a GET request to the target with a 10-second timeout
            requests.post(  # Forward the scraped metrics to Prometheus's PushGateway
                "http://prometheus-pushgateway:9091/metrics/job/batch_scraper",  # PushGateway URL with job label "batch_scraper"
                data=resp.text,  # Use the raw response body as the POST payload (Prometheus text format)
            )
        except Exception as e:  # Catch any exception (network error, timeout, connection refused)
            print(f"Failed to scrape {target}: {e}")  # Log the failure with the target URL and error message

with DAG(  # Enter the DAG context manager to define the workflow structure
    "prometheus-metric-scraper",  # Unique DAG identifier displayed in the Airflow UI and stored in the database
    schedule="*/5 * * * *",  # Cron expression: execute this DAG every 5 minutes
    start_date=datetime(2026, 1, 1),  # DAG start date (prevents automatic backfilling before this date)
    catchup=False,  # Skip creating DAG runs for past schedule intervals that were missed
    tags=["monitoring"],  # Tag this DAG as "monitoring" for easier filtering in the Airflow UI
) as dag:  # Bind the DAG object to the variable `dag`
    scrape = PythonOperator(  # Instantiate a PythonOperator task
        task_id="scrape-internal-metrics",  # Unique task identifier within this DAG
        python_callable=scrape_metrics,  # Reference the Python function to execute when this task runs
    )

from airflow.sensors.http import HttpSensor  # Import the HTTP sensor that checks an external HTTP endpoint

check_model_health = HttpSensor(  # Create an HttpSensor task instance
    task_id="check-prometheus-alert",  # Unique task identifier for Airflow's internal tracking
    http_conn_id="prometheus_conn",  # Reference to an Airflow connection storing Prometheus base URL and credentials
    endpoint="/api/v1/query",  # Prometheus API endpoint for running instant queries
    request_params={  # Dictionary of query parameters appended to the endpoint URL
        "query": 'http_requests_total{endpoint="/invocations", status="5xx"} > 100'  # PromQL: trigger if 5xx errors exceed 100
    },
    response_check=lambda resp: len(resp.json()["data"]["result"]) == 0,  # Return True (healthy) if no results match
    mode="reschedule",  # If the sensor condition is not met, reschedule the check instead of blocking a worker slot
    timeout=300,  # Maximum number of seconds the sensor will wait before timing out
)

from airflow import DAG  # Import the core DAG class for defining Airflow workflows
from airflow.operators.python import PythonOperator  # Import the operator that runs arbitrary Python functions
from datetime import datetime  # Import datetime for specifying the DAG's start_date parameter
import requests  # Import the HTTP library for communicating with Prometheus and Alertmanager APIs

PROMETHEUS = "http://prometheus:9090"  # Internal cluster URL for the Prometheus server
ALERTMANAGER = "http://alertmanager:9093"  # Internal cluster URL for the Alertmanager service

def check_and_alert():  # Define the function that queries Prometheus and fires alerts via Alertmanager
    query = 'rate(cpu_usage{app="ml-model"}[5m]) > 0.8'  # PromQL: trigger if 5-minute average CPU exceeds 80%
    resp = requests.get(  # Send a GET request to Prometheus's instant query API
        f"{PROMETHEUS}/api/v1/query",  # Construct the full URL by appending the API path to the Prometheus base URL
        params={"query": query},  # Pass the PromQL query as a URL query parameter
    )
    results = resp.json()["data"]["result"]  # Extract the result array from Prometheus's JSON response

    if results:  # If the result array is non-empty, the CPU threshold was breached on at least one pod
        alert = {  # Build an alert payload that conforms to Alertmanager's API schema
            "labels": {"alertname": "HighCPU", "severity": "warning"},  # Standard labels: alert name and severity level
            "annotations": {"summary": f"CPU spike detected on {len(results)} pods"},  # Human-readable summary annotation
        }
        requests.post(f"{ALERTMANAGER}/api/v1/alerts", json=[alert])  # POST the alert payload to Alertmanager's API
        print(f"Alert sent: {alert}")  # Log the alert details so they appear in Airflow task logs

with DAG(  # Enter the DAG context manager to define the workflow structure
    "prometheus-self-check",  # Unique DAG identifier displayed in the Airflow UI and stored in the database
    schedule="*/1 * * * *",  # Cron expression: execute this DAG every minute
    start_date=datetime(2026, 1, 1),  # DAG start date (prevents automatic backfilling before this date)
    catchup=False,  # Skip creating DAG runs for past schedule intervals that were missed
) as dag:  # Bind the DAG object to the variable `dag`
    check = PythonOperator(  # Instantiate a PythonOperator task
        task_id="check-prometheus-alerts",  # Unique task identifier within this DAG
        python_callable=check_and_alert,  # Reference the Python function to execute when this task runs
    )