Update Metric collection (#16)

ecmwf · Nov 23, 2023 · 23f71e1 · 23f71e1
1 parent bf9d082
commit 23f71e1
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 16 deletions.
diff --git a/polytope_server/common/metric_collector/queue_metric_collector.py b/polytope_server/common/metric_collector/queue_metric_collector.py
@@ -49,9 +49,18 @@ def total_queued(self):
 
 
 class SQSQueueMetricCollector(QueueMetricCollector):
-    def __init__(self, host):
+    def __init__(self, host, client):
         self.host = host
-        self.message_counts = None
+        self.client = client
 
     def total_queued(self):
-        return sum(self.message_counts.values())
+        response = self.client.get_queue_attributes(
+            QueueUrl=self.host,
+            AttributeNames=[
+                "ApproximateNumberOfMessages",
+                "ApproximateNumberOfMessagesDelayed",
+                "ApproximateNumberOfMessagesNotVisible",
+            ],
+        )
+        values = response.get("Attributes", {}).values()
+        return sum(map(int, values))
diff --git a/polytope_server/common/metric_collector/storage_metric_collector.py b/polytope_server/common/metric_collector/storage_metric_collector.py
@@ -107,12 +107,12 @@ def collect(self):
 
     def storage_space_used(self):
         space_used = 0
-        for db in self.client.list_database_names():
-            space_used += int(getattr(self.client, db).command({"dbStats": 1}).get("storageSize"))
+        # for db in self.client.list_database_names():
+        #     space_used += int(getattr(self.client, db).command({"dbStats": 1}).get("storageSize"))
         return space_used
 
     def total_entries(self):
-        return self.store.count()
+        return self.store.count_documents({})
 
     def db_name(self):
         return self.database

diff --git a/polytope_server/common/queue/sqs_queue.py b/polytope_server/common/queue/sqs_queue.py
@@ -19,7 +19,7 @@ def __init__(self, config):
         self.client = boto3.client("sqs", region_name=region)
         self.queue_url = self.client.get_queue_url(QueueName=queue_name).get("QueueUrl")
         self.check_connection()
-        self.queue_metric_collector = SQSQueueMetricCollector(self.queue_url)
+        self.queue_metric_collector = SQSQueueMetricCollector(self.queue_url, self.client)
 
     def enqueue(self, message):
         self.client.send_message(
@@ -78,13 +78,4 @@ def get_type(self):
         return "sqs"
 
     def collect_metric_info(self):
-        response = self.client.get_queue_attributes(
-            QueueUrl=self.queue_url,
-            AttributeNames=[
-                "ApproximateNumberOfMessages",
-                "ApproximateNumberOfMessagesDelayed",
-                "ApproximateNumberOfMessagesNotVisible",
-            ],
-        )
-        self.queue_metric_collector.message_counts = response["Attributes"]
         return self.queue_metric_collector.collect().serialize()