2222prometheus_username = os .environ .get ("PROMETHEUS_USERNAME" , "" )
2323prometheus_password = os .environ .get ("PROMETHEUS_PASSWORD" , "" )
2424
25+
2526class Prometheus (LoggingConfigurable ):
2627 def query (self , query : str , date_range : DateRange , step : str ) -> requests .Response :
2728 """
@@ -54,13 +55,14 @@ def query(self, query: str, date_range: DateRange, step: str) -> requests.Respon
5455 "step" : step ,
5556 }
5657 query_api = URL (prometheus_api .with_path ("/api/v1/query_range" ))
57- with requests .get (query_api , params = parameters , auth = prometheus_auth ) as response :
58+ with requests .get (
59+ query_api , params = parameters , auth = prometheus_auth
60+ ) as response :
5861 logger .info (f"Querying Prometheus: { response .url } " )
5962 response .raise_for_status ()
6063 result = response .json ()
6164 return result
6265
63-
6466 def query_usage (
6567 self ,
6668 date_range : DateRange ,
@@ -110,7 +112,6 @@ def query_usage(
110112 result = self ._filter_json (result , hub = hub_name , user = user_name )
111113 return result
112114
113-
114115 def _process_response (
115116 self ,
116117 response : requests .Response ,
@@ -149,23 +150,23 @@ def _process_response(
149150 for entry in processed_result :
150151 if "shared" not in entry ["user" ]:
151152 try :
152- entry ["user" ] = escapism .unescape (entry ["user" ], escape_char = "-" )
153+ entry ["user" ] = escapism .unescape (
154+ entry ["user" ], escape_char = "-"
155+ )
153156 except ValueError :
154157 logger .warning (
155158 f"Could not unescape username { entry ['user' ]} for home storage component."
156159 )
157160 continue
158161 return processed_result
159162
160-
161163 def _filter_json (self , result : list [dict ], ** filters ):
162164 return [
163165 item
164166 for item in result
165167 if all (filters [k ] is None or item .get (k ) == filters [k ] for k in filters )
166168 ]
167169
168-
169170 def _pivot_response_dict (self , result : list [dict ]) -> list [dict ]:
170171 """
171172 Pivot the response dictionary to have top-level keys as dates.
@@ -184,7 +185,6 @@ def _pivot_response_dict(self, result: list[dict]) -> list[dict]:
184185 )
185186 return pivot
186187
187-
188188 def _sum_absolute_usage_by_date (self , result : list [dict ]) -> list [dict ]:
189189 """
190190 Sum the absolute usage values by date.
@@ -214,7 +214,6 @@ def _sum_absolute_usage_by_date(self, result: list[dict]) -> list[dict]:
214214 for (date , user , hub , component ), total in sums .items ()
215215 ]
216216
217-
218217 def _calculate_daily_cost_factors (
219218 self , result : list [dict ], hub_name : str | None = None
220219 ) -> list [dict ]:
@@ -256,7 +255,6 @@ def _calculate_daily_cost_factors(
256255 entry ["value" ] = 0.0
257256 return result
258257
259-
260258 @ttl_lru_cache (seconds_to_live = 3600 )
261259 def query_user_groups (
262260 self ,
@@ -277,7 +275,6 @@ def query_user_groups(
277275 result = self ._process_user_groups (response , hub_name , user_name , group_name )
278276 return result
279277
280-
281278 def _process_user_groups (
282279 self ,
283280 response : requests .Response ,
@@ -308,7 +305,6 @@ def _process_user_groups(
308305 )
309306 return result
310307
311-
312308 @ttl_lru_cache (seconds_to_live = 3600 )
313309 def query_users_with_multiple_groups (
314310 self ,
@@ -322,7 +318,12 @@ def query_users_with_multiple_groups(
322318 logger .exception (f"HTTP request failed: { e } " )
323319 raise
324320 grouped = defaultdict (
325- lambda : {"username" : None , "hub" : None , "usergroups" : [], "has_multiple" : False }
321+ lambda : {
322+ "username" : None ,
323+ "hub" : None ,
324+ "usergroups" : [],
325+ "has_multiple" : False ,
326+ }
326327 )
327328 for entry in response :
328329 k = (entry ["username" ], entry ["hub" ])
@@ -343,7 +344,6 @@ def query_users_with_multiple_groups(
343344
344345 return result
345346
346-
347347 @ttl_lru_cache (seconds_to_live = 3600 )
348348 def query_users_with_no_groups (
349349 self ,
0 commit comments