prometheus_metrics_management.py

"""
Alternative approach for working with prometheus metrics.
If the component actively works with metrics,
it is convenient to initialize the controller with labels and further
use methods to change the metric through this controller.

Example:
    # declare controller
    class MetricsController(BaseMetricsController):
        _mandatory_label_names = ['label1', 'label2']
        _optional_label_names = ['label3']

    ...

    # init controller
    controller = MetricsController(labels_map={'label1': 'user-0', 'label2': 'browser-0'})
    # increment value for METRIC
    controller.inc(METRIC, value=10)

    ...

    # temporarily extend labels
    with controller(label3='debug') as ctl:
        # set value for METRIC
        ctl.set(METRIC, value=0)
"""

import json
from abc import ABC
from typing import Any, AnyStr, Dict, List, Protocol, Self, TypeVar, Union

MetricValue = Union[int, float]
MetricLabel = Union[AnyStr, int, float, None]


class Metric(Protocol):
    def labels(self, *args: MetricLabel, **kwargs: MetricLabel) -> Self:
        ...


class SetMetric(Metric, Protocol):
    def set(self, _: MetricValue) -> None:
        ...


class IncrementalMetric(Metric, Protocol):
    def inc(self, _: MetricValue) -> None:
        ...


class BidirectionalMetric(IncrementalMetric, Protocol):
    def dec(self, _: MetricValue) -> None:
        ...


class ObservableMetric(Metric, Protocol):
    def observe(self, _: MetricValue) -> None:
        ...


MetricInstance = TypeVar("MetricInstance", bound=Metric)


class AbstractMetricsController(ABC):
    _mandatory_label_names: List[str]
    _optional_label_names: List[str]
    _empty_label_value: Any = ""

    def __init__(self, labels_map: Dict[str, MetricLabel]) -> None:
        self._validate_labels(labels_map=labels_map)
        self._labels_map = labels_map
        self._extended_labels_map: Dict[str, MetricLabel] = {}

    def __call__(self, **labels_map: MetricLabel) -> Self:
        self._validate_labels(labels_map=labels_map, optional=True)
        self._extended_labels_map.update(labels_map)
        return self

    def __enter__(self) -> Self:
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self._extended_labels_map.clear()

    @classmethod
    def _validate_labels(
        cls, labels_map: Dict[str, MetricLabel], optional: bool = False
    ) -> None:
        check_list = (
            cls._mandatory_label_names if not optional else cls._optional_label_names
        )
        for name, value in labels_map.items():
            if name not in check_list:
                raise ValueError(f'Unknown label "{name}" with value: {value}')

    @property
    def labels(self) -> Dict[str, MetricLabel]:
        total = {**self._labels_map, **self._extended_labels_map}

        for label_name in self._optional_label_names:
            if label_name not in total:
                total[label_name] = self._empty_label_value

        return total


class BaseMetricsController(AbstractMetricsController, ABC):
    _mandatory_label_names: List[str]
    _optional_label_names: List[str]

    def inc(self, metric: IncrementalMetric, value: MetricValue = 1) -> None:
        self._labelling(metric).inc(value)

    def dec(self, metric: BidirectionalMetric, value: MetricValue = 1) -> None:
        self._labelling(metric).dec(value)

    def set(self, metric: SetMetric, value: MetricValue) -> None:
        self._labelling(metric).set(value)

    def observe(self, metric: ObservableMetric, value: MetricValue) -> None:
        self._labelling(metric).observe(value)

    def _labelling(self, metric: MetricInstance) -> MetricInstance:
        return metric.labels(**self.labels)


if __name__ == "__main__":
    from prometheus_client import Gauge
    from prometheus_client.samples import Sample

    METRIC = Gauge(
        name="test_metric",
        labelnames=["label1", "label2", "label3"],
        documentation="",
    )

    class MetricsController(BaseMetricsController):
        _mandatory_label_names = ["label1", "label2"]
        _optional_label_names = ["label3"]

    controller = MetricsController(
        labels_map={"label1": "user-0", "label2": "browser-0"}
    )
    controller.inc(METRIC, 10)

    with controller(label3="debug") as ctl:
        ctl.set(METRIC, -4)

    collection = METRIC.collect()
    samples: List[Sample] = list(collection)[0].samples
    print(f"Metric samples: \n{json.dumps(samples, indent=4)}")

    assert samples[0].labels == {
        "label1": "user-0",
        "label2": "browser-0",
        "label3": "",
    }
    assert samples[0].value == 10.0
    assert samples[1].labels == {
        "label1": "user-0",
        "label2": "browser-0",
        "label3": "debug",
    }
    assert samples[1].value == -4.0