Coverage for aiopromql/client.py: 96%

1 import warnings

2 from datetime import datetime

3 from typing import Optional, Union

4

5 import httpx

6

7 from .models.prometheus import PrometheusResponseModel

8

9

10 class PrometheusClientBase:

11 """Base Prometheus client with common utilities."""

12

13 def __init__(self, url: str):

14 self.base_url = url

15

16 def _parse_response(self, response: dict) -> PrometheusResponseModel:

17 """Parse Prometheus JSON response into model."""

18 return PrometheusResponseModel(**response)

19

20

21 class PrometheusSync(PrometheusClientBase):

22 """Synchronous Prometheus client using httpx."""

23

24 def __init__(self, url: str, timeout: Optional[float] = 2.0):

25 super().__init__(url)

26 self.session = httpx.Client(timeout=httpx.Timeout(timeout))

27

28 def query(self, promql: str, raw: bool = False) -> Union[PrometheusResponseModel, dict]:

29 """

30 Run an instant PromQL query.

31

32 :param promql: The PromQL query string to execute.

33 :type promql: str

34 :param raw: If True, return raw JSON response as dict; otherwise parse into model.

35 :type raw: bool

36 :return: Parsed Prometheus response or raw JSON dict.

37 """

38 response = self.session.get(f"{self.base_url}/api/v1/query", params={"query": promql})

39 response.raise_for_status()

40 return response.json() if raw else self._parse_response(response.json())

41

42 def query_range(

43 self,

44 promql: str,

45 start: datetime,

46 end: datetime,

47 step: str = "30s",

48 raw: bool = False,

49 ) -> Union[PrometheusResponseModel, dict]:

50 """

51 Run a ranged PromQL query over a time window.

52

53 :param promql: The PromQL query string to execute.

54 :type promql: str

55 :param start: Start datetime of the query range.

56 :type start: datetime.datetime

57 :param end: End datetime of the query range.

58 :type end: datetime.datetime

59 :param step: Query resolution step width (e.g., '30s', '1m').

60 :type step: str

61 :param raw: If True, return raw JSON response as dict; otherwise parse into model.

62 :type raw: bool

63 :return: Parsed Prometheus response or raw JSON dict.

64 """

65 start_ts = start.timestamp()

66 end_ts = end.timestamp()

67 response = self.session.get(

68 f"{self.base_url}/api/v1/query_range",

69 params={"query": promql, "start": start_ts, "end": end_ts, "step": step},

70 )

71 response.raise_for_status()

72 return response.json() if raw else self._parse_response(response.json())

73

74 def close(self):

75 """Close the sync client session."""

76 self.session.close()

77

78 def __del__(self):

79 self.close()

80

81

82 class PrometheusAsync(PrometheusClientBase):

83 """Asynchronous Prometheus client using httpx."""

84

85 def __init__(self, url: str, timeout: Optional[float] = 2.0):

86 super().__init__(url)

87 self.client = httpx.AsyncClient(base_url=url, timeout=httpx.Timeout(timeout))

88

89 async def query(self, promql: str, raw: bool = False) -> Union[PrometheusResponseModel, dict]:

90 """

91 Run an instant PromQL query.

92

93 :param promql: The PromQL query string to execute.

94 :type promql: str

95 :param raw: If True, return raw JSON response as dict; otherwise parse into model.

96 :type raw: bool

97 :return: Parsed Prometheus response or raw JSON dict.

98 """

99 response = await self.client.get("/api/v1/query", params={"query": promql})

100 response.raise_for_status()

101 return response.json() if raw else self._parse_response(response.json())

102

103 async def query_range(

104 self,

105 promql: str,

106 start: datetime,

107 end: datetime,

108 step: str = "30s",

109 raw: bool = False,

110 ) -> Union[PrometheusResponseModel, dict]:

111 """

112 Run a ranged PromQL query over a time window.

113

114 :param promql: The PromQL query string to execute.

115 :type promql: str

116 :param start: Start datetime of the query range.

117 :type start: datetime.datetime

118 :param end: End datetime of the query range.

119 :type end: datetime.datetime

120 :param step: Query resolution step width (e.g., '30s', '1m').

121 :type step: str

122 :param raw: If True, return raw JSON response as dict; otherwise parse into model.

123 :type raw: bool

124 :return: Parsed Prometheus response or raw JSON dict.

125 """

126 start_ts = start.timestamp()

127 end_ts = end.timestamp()

128 response = await self.client.get(

129 "/api/v1/query_range",

130 params={"query": promql, "start": start_ts, "end": end_ts, "step": step},

131 )

132 response.raise_for_status()

133 return response.json() if raw else self._parse_response(response.json())

134

135 async def aclose(self):

136 """Close the async client session."""

137 await self.client.aclose()

138

139 async def __aenter__(self):

140 return self

141

142 async def __aexit__(self, exc_type, exc, tb):

143 await self.aclose()

144

145 def __del__(self):

146 if not self.client.is_closed: 146 ↛ 147line 146 didn't jump to line 147 because the condition on line 146 was never true

147 warnings.warn("PrometheusAsync was not closed. Use 'async with' or call 'await .aclose()'")