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

1""" 

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

3""" 

4 

5from collections import defaultdict 

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

7 

8from pydantic import BaseModel 

9 

10from .core import MetricLabelSet, TimeSeries, TimeSeriesPoint 

11 

12 

13class VectorResultModel(BaseModel): 

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

15 

16 metric: Dict[str, str] 

17 value: Tuple[float, str] 

18 

19 

20class MatrixResultModel(BaseModel): 

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

22 

23 metric: Dict[str, str] 

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

25 

26 

27class 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 

50class 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 

74class 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()