cx-ui/backend/triagelib/prometheus.py

"""Prometheus access.

The alert Prometheus lives on the internal 10.11/8 network, which is reachable
only from inside the CX-Tools VPN containers - the laptop itself routes 10.11.*
out of its default gateway. So queries go the same way CX-Tools reaches
OpenStack: `docker exec <region>-osc curl ...`. A direct HTTP transport is tried
first so this still works from a host that does have a route.
"""
from __future__ import annotations

import json
import os
import re
import shlex
import subprocess
import threading
import time
import urllib.parse
import urllib.request
from typing import Any, Optional

DEFAULT_BASE = os.environ.get("CX_PROMETHEUS_BASE", "http://10.11.254.250:9090")

# Containers to try as an HTTP relay, in order. These are the CX-Tools
# OpenStack client containers, which share the regional VPN network namespace.
RELAY_CONTAINERS = ("ca1-osc", "us1-osc", "no1-osc", "ca2-osc")


class PrometheusError(RuntimeError):
    pass


class PrometheusClient:
    def __init__(self, base: str = DEFAULT_BASE, timeout: int = 20):
        self.base = base.rstrip("/")
        self.timeout = timeout
        self._transport: Optional[tuple[str, str]] = None
        self._lock = threading.Lock()

    # --- transport selection ------------------------------------------------

    def _try_direct(self) -> bool:
        try:
            req = urllib.request.Request(f"{self.base}/api/v1/status/buildinfo", headers={"User-Agent": "cx-triage"})
            with urllib.request.urlopen(req, timeout=5) as resp:
                return 200 <= getattr(resp, "status", 200) < 300
        except Exception:
            return False

    def _try_relay(self, container: str) -> bool:
        rc, out, _err = _run(
            ["docker", "exec", "-i", container, "curl", "-sS", "-m", "6", f"{self.base}/api/v1/status/buildinfo"],
            timeout=15,
        )
        return rc == 0 and '"status":"success"' in out

    def transport(self) -> tuple[str, str]:
        """Return (kind, detail) where kind is 'direct' or 'relay'."""
        with self._lock:
            if self._transport is not None:
                return self._transport
            forced = os.environ.get("CX_PROMETHEUS_RELAY", "").strip()
            if forced:
                self._transport = ("relay", forced)
                return self._transport
            if self._try_direct():
                self._transport = ("direct", "host")
                return self._transport
            for container in RELAY_CONTAINERS:
                if self._try_relay(container):
                    self._transport = ("relay", container)
                    return self._transport
            raise PrometheusError(
                f"Cannot reach Prometheus at {self.base}. The host has no route to the internal "
                f"network and none of {', '.join(RELAY_CONTAINERS)} answered. Start the CX-Tools "
                "VPN/OSC containers, or set CX_PROMETHEUS_RELAY to a container that has a route."
            )

    def describe_transport(self) -> str:
        try:
            kind, detail = self.transport()
        except PrometheusError as exc:
            return f"unavailable ({exc})"
        return "direct from host" if kind == "direct" else f"relayed through {detail}"

    # --- requests -----------------------------------------------------------

    def _get(self, path: str, params: Optional[dict[str, str]] = None) -> Any:
        url = f"{self.base}{path}"
        if params:
            url = f"{url}?{urllib.parse.urlencode(params)}"
        kind, detail = self.transport()
        if kind == "direct":
            req = urllib.request.Request(url, headers={"User-Agent": "cx-triage"})
            with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                body = resp.read().decode("utf-8", errors="replace")
        else:
            rc, out, err = _run(
                ["docker", "exec", "-i", detail, "curl", "-sS", "-m", str(self.timeout), url],
                timeout=self.timeout + 10,
            )
            if rc != 0:
                raise PrometheusError(f"Prometheus relay via {detail} failed: {(err or out).strip()}")
            body = out
        try:
            data = json.loads(body)
        except json.JSONDecodeError as exc:
            raise PrometheusError(f"Prometheus returned non-JSON for {path}: {body[:200]}") from exc
        if data.get("status") != "success":
            raise PrometheusError(f"Prometheus error for {path}: {data.get('error') or data}")
        return data.get("data")

    def alerts(self) -> list[dict[str, Any]]:
        data = self._get("/api/v1/alerts") or {}
        alerts = data.get("alerts")
        return [a for a in alerts if isinstance(a, dict)] if isinstance(alerts, list) else []

    def query(self, expr: str) -> list[dict[str, Any]]:
        data = self._get("/api/v1/query", {"query": expr}) or {}
        result = data.get("result")
        return [r for r in result if isinstance(r, dict)] if isinstance(result, list) else []

    def resources_by_floating_ip(self, floating_ip: str) -> list[dict[str, Any]]:
        """The `Resources{floating_ip="..."}` query the Duplicated IPs runbook uses.

        Unlike CX-Tools (production Infrahub only), this series covers every
        environment, so it is how a PreProd/Staging claimant gets found.
        """
        expr = 'Resources{floating_ip="%s"}' % floating_ip.replace('"', "")
        return [dict(r.get("metric") or {}) for r in self.query(expr)]

    def query_range(self, expr: str, start: int, end: int, step: int) -> list[dict[str, Any]]:
        data = self._get("/api/v1/query_range", {
            "query": expr, "start": str(start), "end": str(end), "step": str(step),
        }) or {}
        result = data.get("result")
        return [r for r in result if isinstance(r, dict)] if isinstance(result, list) else []

    def rules(self) -> list[dict[str, Any]]:
        data = self._get("/api/v1/rules") or {}
        return [g for g in (data.get("groups") or []) if isinstance(g, dict)]

    def series_count(self, metric: str) -> int:
        rows = self.query(f"count({metric})")
        if not rows:
            return 0
        try:
            return int(float(rows[0]["value"][1]))
        except (KeyError, IndexError, ValueError, TypeError):
            return 0


class RuleIndex:
    """Maps alertname -> which rule file it came from and its `for` duration.

    Keying off the rule file (not the alert name) is what lets node-exporter
    alerts be separated reliably: two different files both use the group name
    "Imported Rules".
    """

    def __init__(self, client: PrometheusClient, ttl: float = 600.0):
        self.client = client
        self.ttl = ttl
        self._at = 0.0
        self._by_name: dict[str, dict[str, Any]] = {}
        self._error = ""
        self._lock = threading.Lock()

    def refresh(self) -> None:
        groups = self.client.rules()
        index: dict[str, dict[str, Any]] = {}
        for group in groups:
            source = str(group.get("file") or "").rsplit("/", 1)[-1]
            for rule in group.get("rules") or []:
                if rule.get("type") != "alerting":
                    continue
                index[str(rule.get("name") or "")] = {
                    "group": str(group.get("name") or ""),
                    "file": source,
                    "for_seconds": int(rule.get("duration") or 0),
                    "query": str(rule.get("query") or ""),
                }
        self._by_name = index

    def ensure(self) -> None:
        """Refresh if the cache is empty or stale."""
        with self._lock:
            if not self._by_name or time.monotonic() - self._at > self.ttl:
                try:
                    self.refresh()
                    self._error = ""
                except PrometheusError as exc:
                    self._error = str(exc)
                self._at = time.monotonic()

    def get(self, alertname: str) -> dict[str, Any]:
        self.ensure()
        return self._by_name.get(alertname, {})

    @property
    def count(self) -> int:
        return len(self._by_name)

    @property
    def error(self) -> str:
        return self._error


# Metrics the alert rules are built on. If one of these is empty, rules that
# depend on it are broken rather than quiet - see StateSnapshot.broken_inputs.
RULE_INPUT_METRICS = ("Resources", "In_Use_Gpus", "Total_Gpus", "openstack_nova_server_status")

ROGUE_DELTA_EXPR = (
    'sum by (instance) (In_Use_Gpus) - sum by (instance) '
    '(Resources{organization!="3491 - luis.sarabando+runpod@nexgencloud.coms-Organization",'
    'status=~"ACTIVE|SHUTOFF|PRE_ACTIVE"})'
)


def _episode(points: list[tuple[int, float]], median: float) -> dict[str, Any]:
    return {
        "start": points[0][0],
        "end": points[-1][0],
        "minutes": max(1, (points[-1][0] - points[0][0]) // 60 + 1),
        "low": int(min(v for _, v in points)),
        "normal": int(median),
    }


# Labels that identify one alert across time, for true-age lookup.
TRUE_AGE_KEY_LABELS = ("alertname", "instance_name", "floating_ip", "instance", "openstack_id")


def true_age_key(labels: dict[str, Any]) -> tuple:
    """Identity of an alert, built from raw label values on both sides."""
    return tuple(str((labels or {}).get(k, "")) for k in TRUE_AGE_KEY_LABELS)


class TrueAgeIndex:
    """How long each alert's condition has *actually* held.

    Prometheus resets an alert's activeAt whenever the alert resolves, and the
    Infrahub metric pipeline drops most of the `Resources` series for a few
    minutes several times a day. Every alert alive during such a dip resolves and
    re-fires, so activeAt collapses to "time since the last dip" and every alert
    reports the same age.

    This walks the `ALERTS` series backwards instead, bridging gaps shorter than
    GAP_TOLERANCE, which recovers the real duration. Both pending and firing are
    counted, so rules with a long `for:` are not reported as young.
    """

    WINDOW_DAYS = 7
    STEP_SECONDS = 900
    GAP_TOLERANCE = 2700  # 45 min: bridges pipeline dips, not genuine recoveries

    def __init__(self, client: PrometheusClient, severity: str = "infrahub-critical", ttl: float = 300.0):
        self.client = client
        self.severity = severity
        self.ttl = ttl
        self._at = 0.0
        self._lock = threading.Lock()
        self._starts: dict[tuple, tuple[int, bool]] = {}
        self.error = ""
        self.window_start = 0

    def refresh(self) -> None:
        end = int(time.time())
        start = end - self.WINDOW_DAYS * 86400
        self.window_start = start
        expr = (
            "count by (%s) (ALERTS{severity=\"%s\"})"
            % (", ".join(TRUE_AGE_KEY_LABELS), self.severity)
        )
        starts: dict[tuple, tuple[int, bool]] = {}
        for series in self.client.query_range(expr, start, end, self.STEP_SECONDS):
            stamps = []
            for point in series.get("values") or []:
                try:
                    stamps.append(int(float(point[0])))
                except (ValueError, TypeError, IndexError):
                    continue
            if not stamps:
                continue
            run_start = stamps[-1]
            for earlier, later in list(zip(stamps, stamps[1:]))[::-1]:
                if later - earlier > self.GAP_TOLERANCE:
                    break
                run_start = earlier
            # A run that reaches the window edge is only a lower bound.
            capped = run_start <= start + self.STEP_SECONDS
            key = true_age_key(series.get("metric") or {})
            existing = starts.get(key)
            if existing is None or run_start < existing[0]:
                starts[key] = (run_start, capped)
        self._starts = starts

    def get(self) -> "TrueAgeIndex":
        with self._lock:
            if not self._at or time.monotonic() - self._at > self.ttl:
                try:
                    self.refresh()
                    self.error = ""
                except PrometheusError as exc:
                    self.error = str(exc)
                self._at = time.monotonic()
        return self

    def lookup(self, labels: dict[str, Any]) -> tuple[Optional[int], bool]:
        """Return (minutes the condition has held, whether that is a floor)."""
        entry = self._starts.get(true_age_key(labels))
        if not entry:
            return None, False
        start, capped = entry
        return max(0, int((time.time() - start) // 60)), capped

    @property
    def loaded(self) -> bool:
        return bool(self._starts)

    @property
    def count(self) -> int:
        return len(self._starts)


class StateSnapshot:
    """A bulk read of current platform state, used to screen alerts cheaply.

    Re-checking whether an alert's condition still holds is what separates a
    real alert from one that already self-resolved. Doing it from these few
    aggregate queries costs one Prometheus round trip for the whole queue,
    instead of an Infrahub and OpenStack call per alert.
    """

    def __init__(self, client: PrometheusClient, ttl: float = 60.0):
        self.client = client
        self.ttl = ttl
        self._at = 0.0
        self._lock = threading.Lock()
        self.error = ""
        self.by_openstack_id: dict[str, dict[str, str]] = {}
        self.by_instance_name: dict[str, dict[str, str]] = {}
        self.fip_counts: dict[str, int] = {}
        self.rogue_delta: dict[str, float] = {}
        self.total_gpus: dict[str, float] = {}
        self.in_use_gpus: dict[str, float] = {}
        self.resources_by_host: dict[str, list[dict[str, str]]] = {}
        self.broken_inputs: list[str] = []
        # Infrahub VMs in a GPU-counted state with no host recorded. These are
        # invisible to the per-host GPU sum the Rogue VM rule uses, so they can
        # manufacture a gap on whichever host is actually running them.
        self.unattributed_active: int = 0
        self.unattributed_active_gpus: int = 0
        # Episodes where the Resources series partially collapsed. Each one
        # resets activeAt on every alert that was live at the time.
        self.pipeline_dips: list[dict[str, Any]] = []

    def refresh(self) -> None:
        by_osid: dict[str, dict[str, str]] = {}
        by_name: dict[str, dict[str, str]] = {}
        fips: dict[str, int] = {}
        by_host: dict[str, list[dict[str, str]]] = {}
        counted_states = {"ACTIVE", "SHUTOFF", "PRE_ACTIVE"}
        unattributed = 0
        unattributed_gpus = 0

        for row in self.client.query("Resources"):
            metric = {str(k): str(v) for k, v in (row.get("metric") or {}).items()}
            try:
                metric["_gpus"] = str(int(float(row.get("value", [0, "0"])[1])))
            except (ValueError, TypeError, IndexError):
                metric["_gpus"] = "0"
            osid = metric.get("openstack_id", "")
            if osid and osid not in ("None", ""):
                by_osid[osid] = metric
            name = metric.get("instance_name", "")
            if name:
                by_name[name] = metric
            fip = metric.get("floating_ip", "")
            if fip and fip not in ("None", "NULL", ""):
                fips[fip] = fips.get(fip, 0) + 1
            host = metric.get("instance", "")
            if host and host not in ("Unknown", "None"):
                by_host.setdefault(host, []).append(metric)
            elif metric.get("status", "").upper() in counted_states:
                unattributed += 1
                unattributed_gpus += int(metric["_gpus"] or 0)

        self.by_openstack_id = by_osid
        self.by_instance_name = by_name
        self.fip_counts = fips
        self.resources_by_host = by_host
        self.unattributed_active = unattributed
        self.unattributed_active_gpus = unattributed_gpus
        # Summed by instance: these metrics are per (instance, gpu_name), so a
        # host with two GPU models carries two series. Reading them unsummed
        # would silently keep only one.
        self.rogue_delta = self._scalar_by_instance(ROGUE_DELTA_EXPR)
        self.total_gpus = self._scalar_by_instance("sum by (instance) (Total_Gpus)")
        self.in_use_gpus = self._scalar_by_instance("sum by (instance) (In_Use_Gpus)")

        self.broken_inputs = [m for m in RULE_INPUT_METRICS if self.client.series_count(m) == 0]
        self.pipeline_dips = self._find_pipeline_dips()

    def _find_pipeline_dips(self, hours: int = 24, drop_ratio: float = 0.8) -> list[dict[str, Any]]:
        """Find episodes where most of the `Resources` series went missing."""
        end = int(time.time())
        start = end - hours * 3600
        series = self.client.query_range("count(Resources)", start, end, 60)
        if not series:
            return []
        points: list[tuple[int, float]] = []
        for point in series[0].get("values") or []:
            try:
                points.append((int(float(point[0])), float(point[1])))
            except (ValueError, TypeError, IndexError):
                continue
        if len(points) < 10:
            return []
        ordered = sorted(v for _, v in points)
        median = ordered[len(ordered) // 2]
        if median <= 0:
            return []

        episodes: list[dict[str, Any]] = []
        current: list[tuple[int, float]] = []
        for stamp, value in points:
            if value < median * drop_ratio:
                if current and stamp - current[-1][0] > 180:
                    episodes.append(_episode(current, median))
                    current = []
                current.append((stamp, value))
            elif current:
                episodes.append(_episode(current, median))
                current = []
        if current:
            episodes.append(_episode(current, median))
        return episodes

    def _scalar_by_instance(self, expr: str) -> dict[str, float]:
        out: dict[str, float] = {}
        for row in self.client.query(expr):
            host = str((row.get("metric") or {}).get("instance") or "")
            if not host:
                continue
            try:
                out[host] = float(row.get("value", [0, "0"])[1])
            except (ValueError, TypeError, IndexError):
                continue
        return out

    def get(self) -> "StateSnapshot":
        with self._lock:
            if not self._at or time.monotonic() - self._at > self.ttl:
                try:
                    self.refresh()
                    self.error = ""
                except PrometheusError as exc:
                    self.error = str(exc)
                self._at = time.monotonic()
        return self

    @property
    def loaded(self) -> bool:
        return bool(self.by_openstack_id) or bool(self.total_gpus)


def _run(cmd: list[str], timeout: int) -> tuple[int, str, str]:
    try:
        p = subprocess.run(cmd, text=True, capture_output=True, timeout=timeout)
        return p.returncode, p.stdout, p.stderr
    except subprocess.TimeoutExpired as exc:
        return 124, exc.stdout or "", exc.stderr or f"timed out after {timeout}s"
    except FileNotFoundError as exc:
        return 127, "", str(exc)
    except Exception as exc:  # pragma: no cover
        return 1, "", str(exc)


# --- parsing pasted alert text / URLs --------------------------------------

_LABEL_RE = re.compile(r'(\w+)\s*=\s*"((?:[^"\\]|\\.)*)"')


def parse_alert_text(text: str) -> list[dict[str, str]]:
    """Parse pasted `ALERTS{...}` lines, or a Prometheus graph URL, into labels.

    Accepts what a CX engineer actually copies: a raw series line from the
    Prometheus console, or the `g0.expr=` URL from the Slack alert.
    """
    text = (text or "").strip()
    if not text:
        return []

    found: list[dict[str, str]] = []

    # A pasted graph URL: pull the expression out and parse its label matchers.
    if text.startswith("http://") or text.startswith("https://"):
        parsed = urllib.parse.urlparse(text)
        params = urllib.parse.parse_qs(parsed.query)
        exprs = [v for k, vs in params.items() if k.endswith("expr") for v in vs]
        for expr in exprs:
            labels = {k: _unescape(v) for k, v in _LABEL_RE.findall(expr)}
            if labels:
                found.append(labels)
        return found

    for block in re.findall(r"\{[^{}]*\}", text):
        labels = {k: _unescape(v) for k, v in _LABEL_RE.findall(block)}
        if labels:
            found.append(labels)
    if not found:
        labels = {k: _unescape(v) for k, v in _LABEL_RE.findall(text)}
        if labels:
            found.append(labels)
    return found


def _unescape(value: str) -> str:
    return value.replace('\\"', '"').replace("\\\\", "\\").replace("\\n", "\n")


class AlertCache:
    """Short-lived cache so the UI can poll without hammering Prometheus."""

    def __init__(self, client: PrometheusClient, ttl: float = 30.0):
        self.client = client
        self.ttl = ttl
        self._at = 0.0
        self._alerts: list[dict[str, Any]] = []
        self._error = ""
        self._lock = threading.Lock()

    def get(self, force: bool = False) -> tuple[list[dict[str, Any]], str, float]:
        with self._lock:
            age = time.monotonic() - self._at
            if not force and self._at and age < self.ttl:
                return self._alerts, self._error, age
            try:
                self._alerts = self.client.alerts()
                self._error = ""
            except PrometheusError as exc:
                self._error = str(exc)
            self._at = time.monotonic()
            return self._alerts, self._error, 0.0