zfs_metrics.py

#!/usr/bin/env python3

import argparse
import subprocess
import sys
from http.server import BaseHTTPRequestHandler, HTTPServer
from typing import List, Tuple


def run_command(command: List[str]) -> str:
    """Run a command and return stdout as text, raising on errors."""
    return subprocess.check_output(command, text=True, stderr=subprocess.STDOUT)


def parse_zfs_list_output(output: str) -> List[Tuple[str, int, int]]:
    """Parse lines of: name used available (bytes) -> list of tuples."""
    datasets: List[Tuple[str, int, int]] = []
    for raw_line in output.strip().splitlines():
        line = raw_line.strip()
        if not line:
            continue
        parts = line.split()  # -H produces tab-separated, but split on any whitespace is OK
        if len(parts) < 3:
            continue
        name, used_str, avail_str = parts[0], parts[1], parts[2]
        try:
            used_bytes = int(used_str)
            available_bytes = int(avail_str)
        except ValueError:
            continue
        datasets.append((name, used_bytes, available_bytes))
    return datasets


def parse_zpool_list_output(output: str) -> List[Tuple[str, int, int]]:
    """Parse lines of: name allocated free (bytes) -> list of tuples."""
    pools: List[Tuple[str, int, int]] = []
    for raw_line in output.strip().splitlines():
        line = raw_line.strip()
        if not line:
            continue
        parts = line.split()
        if len(parts) < 3:
            continue
        name, allocated_str, free_str = parts[0], parts[1], parts[2]
        try:
            allocated_bytes = int(allocated_str)
            free_bytes = int(free_str)
        except ValueError:
            continue
        pools.append((name, allocated_bytes, free_bytes))
    return pools


def collect_metrics() -> Tuple[List[Tuple[str, int, int]], List[Tuple[str, int, int]], bool, List[str]]:
    """Collect dataset and pool metrics. Returns (datasets, pools, success, errors)."""
    errors: List[str] = []
    success = True

    datasets: List[Tuple[str, int, int]] = []
    pools: List[Tuple[str, int, int]] = []

    try:
        zfs_out = run_command(["zfs", "list", "-Hp", "-o", "name,used,available", "-t", "filesystem,volume"])
        datasets = parse_zfs_list_output(zfs_out)
    except Exception as exc:  # noqa: BLE001
        success = False
        errors.append(f"zfs list failed: {exc}")

    try:
        zpool_out = run_command(["zpool", "list", "-Hp", "-o", "name,allocated,free"])
        pools = parse_zpool_list_output(zpool_out)
    except Exception as exc:  # noqa: BLE001
        success = False
        errors.append(f"zpool list failed: {exc}")

    return datasets, pools, success, errors


def escape_label_value(value: str) -> str:
    """Escape backslash and double-quote per Prometheus exposition format."""
    return value.replace("\\", r"\\").replace('"', r"\"")


def generate_metrics_text() -> str:
    datasets, pools, success, errors = collect_metrics()

    lines: List[str] = []

    # Dataset metrics
    lines.append("# HELP zfs_dataset_used_bytes Used space of ZFS datasets in bytes.")
    lines.append("# TYPE zfs_dataset_used_bytes gauge")
    for name, used_b, _ in datasets:
        lines.append(f'zfs_dataset_used_bytes{{dataset="{escape_label_value(name)}"}} {used_b}')

    lines.append("# HELP zfs_dataset_available_bytes Available space of ZFS datasets in bytes.")
    lines.append("# TYPE zfs_dataset_available_bytes gauge")
    for name, _, avail_b in datasets:
        lines.append(f'zfs_dataset_available_bytes{{dataset="{escape_label_value(name)}"}} {avail_b}')

    # Pool metrics
    lines.append("# HELP zfs_pool_allocated_bytes Allocated (used) space of ZFS pools in bytes.")
    lines.append("# TYPE zfs_pool_allocated_bytes gauge")
    for pool_name, allocated_b, _ in pools:
        lines.append(f'zfs_pool_allocated_bytes{{pool="{escape_label_value(pool_name)}"}} {allocated_b}')

    lines.append("# HELP zfs_pool_free_bytes Free space of ZFS pools in bytes.")
    lines.append("# TYPE zfs_pool_free_bytes gauge")
    for pool_name, _, free_b in pools:
        lines.append(f'zfs_pool_free_bytes{{pool="{escape_label_value(pool_name)}"}} {free_b}')

    # Exporter success metric
    lines.append("# HELP zfs_exporter_success 1 if last scrape succeeded, else 0.")
    lines.append("# TYPE zfs_exporter_success gauge")
    lines.append(f"zfs_exporter_success {1 if success else 0}")

    # Optional error comments (comments are ignored by Prometheus)
    for err in errors:
        lines.append(f"# error: {err}")

    return "\n".join(lines) + "\n"


class MetricsHandler(BaseHTTPRequestHandler):
    def do_GET(self) -> None:  # noqa: N802  (match BaseHTTPRequestHandler signature)
        if self.path not in ("/metrics", "/"):
            self.send_response(404)
            self.end_headers()
            self.wfile.write(b"Not Found\n")
            return

        metrics_text = generate_metrics_text()
        body = metrics_text.encode("utf-8")
        self.send_response(200)
        self.send_header("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, fmt: str, *args) -> None:  # silence default stdout logging
        return


def serve_http(address: str, port: int) -> None:
    server = HTTPServer((address, port), MetricsHandler)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        server.server_close()


def main(argv: List[str]) -> int:
    parser = argparse.ArgumentParser(description="ZFS Prometheus metrics exporter")
    parser.add_argument(
        "--mode",
        choices=["stdout", "http"],
        default="stdout",
        help="stdout: print once and exit; http: serve /metrics",
    )
    parser.add_argument("--address", default="0.0.0.0", help="Listen address for HTTP mode (default: 0.0.0.0)")
    parser.add_argument("--port", type=int, default=9808, help="Listen port for HTTP mode (default: 9808)")

    args = parser.parse_args(argv)

    if args.mode == "stdout":
        sys.stdout.write(generate_metrics_text())
        return 0

    serve_http(args.address, args.port)
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))