FabledSteward/steward/core/prometheus.py

"""Prometheus text exposition format (0.0.4): one parser for every scraper.

The Traefik plugin and the Prometheus plugin both read this format. parse()
returns families keyed by name, each carrying its declared type, so a caller
can tell a counter (store a rate) from a gauge (store the value) without
guessing from the name.

Samples are grouped under the family their TYPE line declares: a histogram
"x" owns x_bucket, x_sum and x_count; a summary owns x, x_sum and x_count.
A sample with no TYPE line is its own untyped family. Timestamps are read
past and ignored; a line that doesn't parse is skipped, never fatal.
"""
from __future__ import annotations

import math
import re
from dataclasses import dataclass, field

_NAME = re.compile(r"[a-zA-Z_:][a-zA-Z0-9_:]*")
_LABEL_NAME = re.compile(r"[a-zA-Z_][a-zA-Z0-9_]*")

TYPES = ("counter", "gauge", "histogram", "summary", "untyped")

# Suffixes a histogram or summary's samples carry beyond the family name.
_CHILD_SUFFIXES = {
    "counter": ("_total",),        # OpenMetrics declares a counter without it
    "histogram": ("_bucket", "_sum", "_count"),
    "summary": ("_sum", "_count"),
}

_ESCAPES = {"\\": "\\", '"': '"', "n": "\n"}


@dataclass(frozen=True)
class Sample:
    name: str
    labels: tuple[tuple[str, str], ...]   # in the order the line gave them
    value: float

    def label_dict(self) -> dict[str, str]:
        return dict(self.labels)


@dataclass
class Family:
    name: str
    type: str = "untyped"
    help: str = ""
    samples: list[Sample] = field(default_factory=list)


def parse(text: str) -> dict[str, Family]:
    """{family name: Family}, in the order families first appear."""
    families: dict[str, Family] = {}
    for raw in text.splitlines():
        line = raw.strip()
        if not line:
            continue
        if line.startswith("#"):
            _read_comment(line, families)
            continue
        sample = parse_sample(line)
        if sample is None:
            continue
        fam = families.get(_owner(sample.name, families))
        if fam is None:
            fam = families[sample.name] = Family(sample.name)
        fam.samples.append(sample)
    return families


def _read_comment(line: str, families: dict[str, Family]) -> None:
    parts = line[1:].strip().split(None, 2)
    if len(parts) < 2 or parts[0] not in ("TYPE", "HELP"):
        return
    name = parts[1]
    if not _NAME.fullmatch(name):
        return
    fam = families.setdefault(name, Family(name))
    rest = parts[2].strip() if len(parts) > 2 else ""
    if parts[0] == "TYPE":
        fam.type = rest if rest in TYPES else "untyped"
    else:
        fam.help = rest.replace("\\n", "\n").replace("\\\\", "\\")


def _owner(sample_name: str, families: dict[str, Family]) -> str:
    """The declared family a sample belongs to: its own name, or the
    histogram/summary whose name it extends."""
    if sample_name in families:
        return sample_name
    for kind, suffixes in _CHILD_SUFFIXES.items():
        for suffix in suffixes:
            if sample_name.endswith(suffix):
                base = sample_name[: -len(suffix)]
                fam = families.get(base)
                if fam is not None and fam.type == kind:
                    return base
    return sample_name


def parse_sample(line: str) -> Sample | None:
    """One sample line, or None when it isn't one."""
    m = _NAME.match(line)
    if not m:
        return None
    name, pos = m.group(), m.end()
    labels: tuple[tuple[str, str], ...] = ()
    if pos < len(line) and line[pos] == "{":
        read = _read_labels(line, pos + 1)
        if read is None:
            return None
        labels, pos = read
    rest = line[pos:].split()
    if not rest or (pos < len(line) and not line[pos].isspace()):
        return None
    value = parse_value(rest[0])
    if value is None:
        return None
    return Sample(name, labels, value)


def _read_labels(line: str, pos: int) -> tuple[tuple[tuple[str, str], ...], int] | None:
    """Labels from just after "{" to just after "}", values unescaped.
    Commas and braces inside a quoted value are part of the value."""
    out: list[tuple[str, str]] = []
    n = len(line)
    while True:
        while pos < n and line[pos] in " \t,":
            pos += 1
        if pos >= n:
            return None
        if line[pos] == "}":
            return tuple(out), pos + 1
        m = _LABEL_NAME.match(line, pos)
        if not m:
            return None
        key, pos = m.group(), m.end()
        while pos < n and line[pos] in " \t":
            pos += 1
        if line[pos:pos + 2] != '="':
            return None
        pos += 2
        chars: list[str] = []
        while True:
            if pos >= n:
                return None
            c = line[pos]
            if c == "\\" and pos + 1 < n:
                nxt = line[pos + 1]
                chars.append(_ESCAPES.get(nxt, "\\" + nxt))
                pos += 2
                continue
            if c == '"':
                pos += 1
                break
            chars.append(c)
            pos += 1
        out.append((key, "".join(chars)))


def parse_value(text: str) -> float | None:
    """A sample value: a float, NaN, +Inf or -Inf (Go's spellings too)."""
    lowered = text.lower()
    if lowered in ("+inf", "inf"):
        return math.inf
    if lowered == "-inf":
        return -math.inf
    if lowered == "nan":
        return math.nan
    try:
        return float(text)
    except ValueError:
        return None