Skip to content

Instantly share code, notes, and snippets.

@maratuska
Created February 21, 2023 09:48
Show Gist options
  • Save maratuska/28e6227a4731260b39e83ae9c16b8bc4 to your computer and use it in GitHub Desktop.
Save maratuska/28e6227a4731260b39e83ae9c16b8bc4 to your computer and use it in GitHub Desktop.
"""
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment