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.