From 1dae9046f2a3e47ae23ecac86f0255be768e51fd Mon Sep 17 00:00:00 2001
From: dettlaff <dettlaff@riseup.net>
Date: Fri, 26 Jul 2024 15:50:12 +0400
Subject: [PATCH] feat: add union type

---
 selfprivacy_api/graphql/queries/monitoring.py | 17 +++++++---
 selfprivacy_api/utils/prometheus.py           | 31 ++++++++++++++-----
 2 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/selfprivacy_api/graphql/queries/monitoring.py b/selfprivacy_api/graphql/queries/monitoring.py
index 637dc2a..9179268 100644
--- a/selfprivacy_api/graphql/queries/monitoring.py
+++ b/selfprivacy_api/graphql/queries/monitoring.py
@@ -1,7 +1,7 @@
 import strawberry
 from typing import Optional
 from datetime import datetime
-from selfprivacy_api.utils.prometheus import PrometheusQueries, PrometheusQueryResult
+from selfprivacy_api.utils.prometheus import PrometheusQueries, Response
 
 
 @strawberry.type
@@ -12,7 +12,7 @@ class Monitoring:
         start: Optional[datetime] = None,
         end: Optional[datetime] = None,
         step: int = 60,
-    ) -> PrometheusQueryResult:
+    ) -> Response:
         return PrometheusQueries.disk_usage(start, end, step)
 
     @strawberry.field
@@ -21,7 +21,7 @@ class Monitoring:
         start: Optional[datetime] = None,
         end: Optional[datetime] = None,
         step: int = 60,
-    ) -> PrometheusQueryResult:
+    ) -> Response:
         return PrometheusQueries.memory_usage(start, end, step)
 
     @strawberry.field
@@ -30,5 +30,14 @@ class Monitoring:
         start: Optional[datetime] = None,
         end: Optional[datetime] = None,
         step: int = 60,
-    ) -> PrometheusQueryResult:
+    ) -> Response:
+        return PrometheusQueries.cpu_usage(start, end, step)
+
+    @strawberry.field
+    def network_usage(
+        self,
+        start: Optional[datetime] = None,
+        end: Optional[datetime] = None,
+        step: int = 60,
+    ) -> Response:
         return PrometheusQueries.cpu_usage(start, end, step)
diff --git a/selfprivacy_api/utils/prometheus.py b/selfprivacy_api/utils/prometheus.py
index a8b4daa..ad071b2 100644
--- a/selfprivacy_api/utils/prometheus.py
+++ b/selfprivacy_api/utils/prometheus.py
@@ -7,7 +7,7 @@ import strawberry
 from strawberry.scalars import JSON
 
 from dataclasses import dataclass
-from typing import Optional
+from typing import Optional, Annotated, Union
 from datetime import datetime, timedelta
 
 PROMETHEUS_URL = "http://localhost:9001"
@@ -20,9 +20,20 @@ class PrometheusQueryResult:
     result: JSON
 
 
+@strawberry.type
+class PrometheusQueryError:
+    error: str
+
+
+Response = Annotated[
+    Union[PrometheusQueryResult, PrometheusQueryError],
+    strawberry.union("PrometheusQueryResponse"),
+]
+
+
 class PrometheusQueries:
     @staticmethod
-    def _send_query(query: str, start: int, end: int, step: int):
+    def _send_query(query: str, start: int, end: int, step: int) -> Response:
         try:
             response = requests.get(
                 f"{PROMETHEUS_URL}/api/v1/query",
@@ -34,20 +45,24 @@ class PrometheusQueries:
                 },
             )
             if response.status_code != 200:
-                raise Exception("Prometheus returned unexpected HTTP status code")
+                return PrometheusQueryError(
+                    error="Prometheus returned unexpected HTTP status code"
+                )
             json = response.json()
             return PrometheusQueryResult(
                 result_type=json["data"]["resultType"], result=json["data"]["result"]
             )
         except Exception as error:
-            raise Exception("Prometheus request failed! " + str(error))
+            return PrometheusQueryError(
+                error=f"Prometheus request failed! Error: {str(error)}"
+            )
 
     @staticmethod
     def cpu_usage(
         start: Optional[datetime] = None,
         end: Optional[datetime] = None,
         step: int = 60,  # seconds
-    ) -> PrometheusQueryResult:
+    ) -> Response:
         """
         Get CPU information.
 
@@ -82,7 +97,7 @@ class PrometheusQueries:
         start: Optional[datetime] = None,
         end: Optional[datetime] = None,
         step: int = 60,  # seconds
-    ) -> PrometheusQueryResult:
+    ) -> Response:
         """
         Get memory usage.
 
@@ -117,7 +132,7 @@ class PrometheusQueries:
         start: Optional[datetime] = None,
         end: Optional[datetime] = None,
         step: int = 60,  # seconds
-    ) -> PrometheusQueryResult:
+    ) -> Response:
         """
         Get disk usage information.
 
@@ -152,7 +167,7 @@ class PrometheusQueries:
         start: Optional[datetime] = None,
         end: Optional[datetime] = None,
         step: int = 60,  # seconds
-    ) -> PrometheusQueryResult:
+    ) -> Response:
         """
         Get network usage information for both download and upload.