Add in used_diskspace to catalog_worker

cedadev · Feb 25, 2025 · b4eb34c · b4eb34c
1 parent c75f9f7
commit b4eb34c
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 3 deletions.
diff --git a/nlds/rabbit/publisher.py b/nlds/rabbit/publisher.py
@@ -113,6 +113,7 @@ class RabbitMQPublisher():
     MSG_USER = "user"
     MSG_GROUP = "group"
     MSG_QUOTA = "quota"
+    MSG_DISKSPACE = "diskspace"
     MSG_GROUPALL = "groupall"
     MSG_TARGET = "target"
     MSG_ROUTE = "route"

diff --git a/nlds_processors/catalog/catalog_worker.py b/nlds_processors/catalog/catalog_worker.py
@@ -1856,7 +1856,7 @@ def _catalog_meta(self, body: Dict, properties: Header) -> None:
         )
 
     def _catalog_quota(self, body: Dict, properties: Header) -> None:
-        """Return the users quota for the given service."""
+        """Return the user's quota for the given group."""
         message_vars = self._parse_user_vars(body)
         if message_vars is None:
             # Check if any problems have occured in the parsing of the message
@@ -1871,18 +1871,37 @@ def _catalog_quota(self, body: Dict, properties: Header) -> None:
         try:
             group_quota = self.authenticator.get_tape_quota(service_name=group)
         except CatalogError as e:
-            # failed to get the holdings - send a return message saying so
+            # failed to get the tape quota - send a return message saying so
             self.log(e.message, self.RK_LOG_ERROR)
             body[self.MSG_DETAILS][self.MSG_FAILURE] = e.message
             body[self.MSG_DATA][self.MSG_QUOTA] = None
         else:
-            # fill the return message with a dictionary of the holding(s)
+            # fill the return message with the group quota
             body[self.MSG_DATA][self.MSG_QUOTA] = group_quota
             self.log(
                 f"Quota from CATALOG_QUOTA {group_quota}",
                 self.RK_LOG_DEBUG
             )
 
+        self.catalog.start_session()
+
+        try:
+            used_diskspace = self.catalog.get_used_diskspace(user=user, group=group)
+        except CatalogError as e:
+            # failed to get the used diskspace - send a return message saying so
+            self.log(e.message, self.RK_LOG_ERROR)
+            body[self.MSG_DETAILS][self.MSG_FAILURE] = e.message
+            body[self.MSG_DATA][self.MSG_DISKSPACE] = None
+        else:
+            # fill the return message with the used diskspace
+            body[self.MSG_DATA][self.MSG_DISKSPACE] = used_diskspace
+            self.log(
+                f"Used diskspace from CATALOG_QUOTA {used_diskspace}",
+                self.RK_LOG_DEBUG
+            )
+
+        self.catalog.end_session()
+
         # return message to complete RPC
         self.publish_message(
             properties.reply_to,