Add PureGym Prometheus stats (#28)

PulumiWebServer/Nix/puregym/puregym.py

@@ -13,12 +13,49 @@ class MyHandler(BaseHTTPRequestHandler):
     _last_accessed_by_id = defaultdict(lambda: datetime.min)
     _last_accessed_by_name = defaultdict(lambda: datetime.min)
 
+    _all_gyms = {}
+    _last_refreshed_gyms = datetime.min
+
     def _bad_request(self, text: str, code: int = 400) -> None:
         self.send_response(code)
         self.send_header('Content-type', 'text/plain; charset=utf-8')
         self.end_headers()
         self.wfile.write(text.encode('utf-8'))
 
+    def _get_fullness(self, gym_id: int) -> bytes:
+        if abs(datetime.now() - self._last_accessed_by_id[gym_id]) > timedelta(seconds=30):
+            token = subprocess.check_output(['cat', '/tmp/puregym_token']).strip()
+            output = subprocess.check_output(
+                [puregym, 'fullness', '--bearer-token', token, '--gym-id', str(gym_id)], text=True,
+                encoding='utf-8')
+            output = output.encode('utf-8')
+            self._cache_result_by_id[gym_id] = output
+            self._last_accessed_by_id[gym_id] = datetime.now()
+        else:
+            output = self._cache_result_by_id[gym_id]
+        return output
+
+    def _refresh_gyms(self) -> None:
+        if self._last_refreshed_gyms < datetime.now() - timedelta(days=1):
+            token = subprocess.check_output(['cat', '/tmp/puregym_token']).strip()
+            output = subprocess.check_output(
+                [puregym, 'all-gyms', '--bearer-token', token, '--terse', 'true'], text=True,
+                encoding='utf-8')
+            new_gyms = {}
+            for line in output.splitlines():
+                gym_id, gym_name = line.split(',')
+                new_gyms[int(gym_id)] = gym_name
+            self._all_gyms = new_gyms
+            self._last_refreshed_gyms = datetime.now()
+
+    def get_all_gyms(self, _query: dict[AnyStr, list[AnyStr]]) -> None:
+        self._refresh_gyms()
+        self.send_response(200)
+        self.send_header('Content-type', 'text/plain')
+        self.end_headers()
+        for gym_id, gym_name in self._all_gyms.items():
+            self.wfile.write(f'{gym_id}: {gym_name}\n'.encode())
+
     def get_fullness(self, query: dict[AnyStr, list[AnyStr]]) -> None:
         desired_gym_name = None
         query_gym = query.get("gym_name", None)
@@ -48,16 +85,7 @@ class MyHandler(BaseHTTPRequestHandler):
             desired_gym_id = None
 
         if desired_gym_id is not None:
-            if abs(datetime.now() - self._last_accessed_by_id[desired_gym_id]) > timedelta(seconds=30):
-                token = subprocess.check_output(['cat', '/tmp/puregym_token']).strip()
-                output = subprocess.check_output(
-                    [puregym, 'fullness', '--bearer-token', token, '--gym-id', str(desired_gym_id)], text=True,
-                    encoding='utf-8')
-                output = output.encode('utf-8')
-                self._cache_result_by_id[desired_gym_id] = output
-                self._last_accessed_by_id[desired_gym_id] = datetime.now()
-            else:
-                output = self._cache_result_by_id[desired_gym_id]
+            output = self._get_fullness(desired_gym_id)
         elif desired_gym_name is not None:
             if abs(datetime.now() - self._last_accessed_by_name[desired_gym_name]) > timedelta(seconds=30):
                 token = subprocess.check_output(['cat', '/tmp/puregym_token']).strip()
@@ -79,8 +107,39 @@ class MyHandler(BaseHTTPRequestHandler):
         self.end_headers()
         self.wfile.write(output)
 
+    def get_prometheus(self, query: dict[AnyStr, list[AnyStr]]) -> None:
+        query_gym = query.get("gym_id", None)
+        if query_gym is None:
+            self._bad_request('Must supply gym_id')
+            return
+        try:
+            gym_id = [int(i) for i in query_gym]
+        except ValueError:
+            self._bad_request('at least one gym_id did not parse as an int')
+            return
+
+        if not gym_id:
+            self._bad_request('supply at least one gym_id')
+            return
+
+        try:
+            fullness = [(i, int(self._get_fullness(i).split(b' ')[0])) for i in gym_id]
+        except ValueError:
+            self._bad_request('at least one fullness did not yield an int', 500)
+            return
+
+        self.send_response(200)
+        self.send_header('Content-type', 'text/plain')
+        self.end_headers()
+        self._refresh_gyms()
+        for gym_id, fullness in fullness:
+            gym_name = ''.join(c for c in self._all_gyms[gym_id] if c == ' ' or str.isalnum(c))
+            self.wfile.write(f'fullness{{label="{gym_name}"}} {fullness}\n'.encode())
+
     _handlers: dict[str, Callable[["MyHandler", dict[AnyStr, list[AnyStr]]], None]] = {
-        "/fullness": get_fullness
+        "/fullness": get_fullness,
+        "/fullness-prometheus": get_prometheus,
+        "/gym-mapping": get_all_gyms,
     }
 
     def do_GET(self):