Coverage for aiopromql/models/prometheus.py: 100%

1"""

2 Pydantic models for parsing and transforming Prometheus query json responses.

3"""

4

5 from collections import defaultdict

6 from typing import Dict, List, Literal, Tuple, Union

7

8 from pydantic import BaseModel

9

10 from .core import MetricLabelSet, TimeSeries, TimeSeriesPoint

11

12

13 class VectorResultModel(BaseModel):

14 """Single Prometheus vector result entry."""

15

16 metric: Dict[str, str]

17 value: Tuple[float, str]

18

19

20 class MatrixResultModel(BaseModel):

21 """Single Prometheus matrix result entry."""

22

23 metric: Dict[str, str]

24 values: List[Tuple[float, str]]

25

26

27 class VectorDataModel(BaseModel):

28 """Parsed vector data block from Prometheus."""

29

30 resultType: Literal["vector"]

31 result: List[VectorResultModel]

32

33 def to_metric_map(self) -> Dict[MetricLabelSet, TimeSeries]:

34 """

35 Converts vector results to a dict of TimeSeries object grouped by metric labels.

36

37 Returns:

38 Dictionary mapping MetricLabelSet to TimeSeries.

39 """

40 metric_map: Dict[MetricLabelSet, TimeSeries] = defaultdict(lambda: TimeSeries([]))

41

42 for r in self.result:

43 key = MetricLabelSet(r.metric)

44 ts_point = TimeSeriesPoint.from_prometheus_value(*r.value)

45 metric_map[key].add_point(ts_point)

46

47 return dict(metric_map)

48

49

50 class MatrixDataModel(BaseModel):

51 """Parsed matrix data block from Prometheus."""

52

53 resultType: Literal["matrix"]

54 result: List[MatrixResultModel]

55

56 def to_metric_map(self) -> Dict[MetricLabelSet, TimeSeries]:

57 """

58 Converts matrix results to a dict of TimeSeries grouped by metric labels.

59

60 Returns:

61 Dictionary mapping MetricLabelSet to TimeSeries.

62 """

63 metric_map: Dict[MetricLabelSet, TimeSeries] = defaultdict(lambda: TimeSeries([]))

64

65 for r in self.result:

66 key = MetricLabelSet(r.metric)

67 for val in r.values:

68 ts_point = TimeSeriesPoint.from_prometheus_value(*val)

69 metric_map[key].add_point(ts_point)

70

71 return dict(metric_map)

72

73

74 class PrometheusResponseModel(BaseModel):

75 """Top-level Prometheus query response wrapper."""

76

77 status: Literal["success"]

78 data: Union[VectorDataModel, MatrixDataModel]

79

80 def to_metric_map(self) -> Dict[MetricLabelSet, TimeSeries]:

81 """

82 Converts the response into a mapping from metric label sets to time series.

83

84 The resulting data structure uses generic, reusable types:

85 - MetricLabelSet: a hashable representation of metric labels.

86 - TimeSeries: a sequence of timestamped values.

87

88 Returns:

89 A dictionary mapping MetricLabelSet to TimeSeries.

90 """

91 return self.data.to_metric_map()