improved quota stats script

scripts/authN_and_authZ/conf/quota.stats.conf

@@ -1,11 +1,13 @@
 {
 "stat_output_file": "accounting/quota.stats.json",
-"log_level": "DEBUG",
+"storage_space_unity": "GB",
+"log_level": "INFO",
 "log_file": "quota.stats.log",
-"notification_sender": "proirod1",
-"notification_receiver": "[email protected]",
+"notification_active": "True",
+"notification_sender": "xxxxxxxx",
+"notification_receiver": "[email protected]",
 "internal_project_list": "public,rodsadmin",
 "mirrored_projects": "test1,test2",
-"irods_home_dir": "/CINECA01/home/",
-"irods_debug": "True"
+"irods_home_dir": "/TestZone/home/",
+"irods_debug": "False"
 }

scripts/authN_and_authZ/quota_stats.py

@@ -17,7 +17,7 @@
 logger = logging.getLogger('QuotaStats')
 
 ################################################################################
-# User Accounts Synchronization Class #
+# Collection Quota Stats Class #
 ################################################################################
 
 class QuotaStats():
@@ -86,24 +86,84 @@ def _writeUsageStats(self,usage):
 
         return True
 
-
-    def _getUsedSpace(self,proj_name):
-        """Internal: get the used space from iRODS iCAT in GB"""
 
-        query = "select sum(DATA_SIZE) where COLL_NAME like '/CINECA01/home/"+ proj_name +"%'"
+    def _getStorageInfo(self, query):
+        """Internal: get the storage info from iRODS iCAT"""
+
         logger.debug("executing the following query: %s", query)
         [rc,out] = self.irodsu.queryIrodsIcat(query)
         if rc != 0:
-            logger.error("Error: %s", err)
-            return False
-        [key,quota] = out.split('=')
-        quotaGB = 0
-        quota = quota.replace("\n------------------------------------------------------------\n","")
-        if len(quota.strip()) > 0:
-            quotaGB = int(quota.strip())/1074000000
+            logger.error("Error: the query '%s' failed", query)
+            return None
+        [key,info] = out.split('=')
+        info = info.replace("\n------------------------------------"
+                            + "------------------------\n","")
+        return info
+
+
+    def _getUsedSpace(self, proj_name):
+        """Internal: get the used space from iRODS iCAT"""
+
+        query = "select sum(DATA_SIZE) where COLL_NAME like '" \
+                + self.conf.irods_home_dir + proj_name +"%'"
+        info = self._getStorageInfo(query)
+        if info is not None:
+            if len(info.strip()) > 0:
+                return self._fromBytes(int(info.strip()), 
+                                       self.conf.storage_space_unity)
+            else:
+                return 0
+        else:
+            return None
+
+
+    def _getTotalNumberObjects(self, proj_name):
+        """Internal: get the total number of object stored under the specified
+           project
+        """
+
+        query = "select count(DATA_ID) where COLL_NAME like '" \
+                + self.conf.irods_home_dir + proj_name + "%'"
+        info = self._getStorageInfo(query)
+        if info is not None:
+            return int(info.strip())
+        else:
+            return None
+
+
+    def _fromBytes(self, size, unity):
+        """Convert file size from byte"""
+        size_map = {'B': 1, 'KB': 1024, 'MB': 1024 ** 2, 'GB': 1024 ** 3,
+                    'TB': 1024 ** 4}
+        return size / size_map[unity]
+
+
+    def _thresholdAlarm(self, project_usage, proj_name, old_used_space_perc):
+        """Notification mechanism to alert about over quota events"""
+
+        message = ""
+        threshold_soft = 95
+        threshold_hard = 100
+        if (proj_name in self.conf.mirrored_projects):
+            threshold_soft = 95*2
+            threshold_hard = 100*2
+        if ((project_usage['used_space_perc'] > threshold_soft) and 
+            (project_usage['used_space_perc'] > old_used_space_perc)):
+            message = "project " + proj_name + " is reaching its quota limit " \
+                    + "(used space > 95%): " + str(project_usage['used_space'])\
+                    + " " + self.conf.storage_space_unity
+        if (project_usage['used_space_perc'] >= threshold_hard):
+            message = "project " + proj_name + " reached its quota limit " \
+                    + "(used space > 100%): " + str(project_usage['used_space'])\
+                    + " " + self.conf.storage_space_unity
+        if (len(message) > 0):
+            mailsnd = MailSender()
+            mailsnd.send(message, self.conf.notification_sender, 
+                         self.conf.notification_receiver)
+            logger.debug("sent alert for quota over limit related to project: "
+                         + proj_name)
+
 
-        return quotaGB
-
     # Public methods
 
     def computeStats(self,dryrun):
@@ -114,59 +174,81 @@ def computeStats(self,dryrun):
         
         """
 
+        logger.info("Getting storage info from file")
         usage = self._parseUsageStats()
         # if log level is equal to DEBUG (level 10)
         if (logger.getEffectiveLevel() == 10):
             pp = pprint.PrettyPrinter(indent=4)
-            print("Usage info:")
+            print("Storage info:")
             pp.pprint(usage)
 
+        total_used_space = 0
+        total_allocated_space = 0
+        total_number_of_objects = 0
         for proj_name in [x for x in self._parseIrodsProj().keys()
                           if x not in self.conf.internal_project_list]:
 
+            logger.info("Looking at project: " + proj_name)
             project_usage = {}
+
+            logger.info("Getting storage info")
             if not(dryrun):
-                quotaGB = self._getUsedSpace(proj_name)
+                quota = self._getUsedSpace(proj_name)
+            elif usage and (proj_name in usage.keys()):
+                logger.info("Running in dryrun mode")
+                quota = usage[proj_name]['used_space']
             elif usage and (proj_name in usage.keys()):
                 quotaGB = usage[proj_name]['used_space']
             else:
-                quotaGB = 0
-            limit = 0
-            old_used_space_perc = 0
-            if usage and (proj_name in usage.keys()):
-                project_usage = usage[proj_name]
-                project_usage['used_space'] = quotaGB
-                limit = int(usage[proj_name]['quota_limit'])
-                old_used_space_perc = usage[proj_name]['used_space_perc']
-            if (len((str(limit)).strip()) > 0 and limit > 0):
-                project_usage['used_space_perc'] = quotaGB*100/limit
-            else:
-                continue
-            #TODO add the silent option to the notification mechanism
-            message = ""
-            threshold_soft = 95
-            threshold_hard = 100
-            if (proj_name in self.conf.mirrored_projects):
-                threshold_soft = 95*2
-                threshold_hard = 100*2
-            if ((project_usage['used_space_perc'] > threshold_soft) and \
-                (project_usage['used_space_perc'] > old_used_space_perc)):
-                message = "project "+proj_name+" is reaching its quota limit (used space > 95%): "\
-                          + str(quotaGB) + " GB"
-            if (project_usage['used_space_perc'] >= threshold_hard):
-                message = "project "+proj_name+" reached its quota limit (used space >= 100%): "\
-                          + str(quotaGB) + " GB"
-            if (len(message) > 0):
-                mailsnd = MailSender()
-                mailsnd.send(message, self.conf.notification_sender, self.conf.notification_receiver)
-                logger.debug("sent alert for quota over limit related to project: "
-                             + proj_name)
-            if usage: usage[proj_name] = project_usage
+                logger.info("No storage information available")
+                quota = 0
+
+            if quota is not None:
+
+                logger.info("Getting the number of objects")
+                if not(dryrun):
+                    num_objs = self._getTotalNumberObjects(proj_name)
+                elif usage and (proj_name in usage.keys()):
+                    logger.info("Running in dryrun mode")
+                    num_objs = usage[proj_name]['number_of_objects']
+                else:
+                    logger.info("No storage information available")
+                    num_objs = 0
+
+                limit = 0
+                old_used_space_perc = 0
+                if usage and (proj_name in usage.keys()):
+                    project_usage = usage[proj_name]
+                    project_usage['used_space'] = quota
+                    project_usage['number_of_objects'] = num_objs
+                    limit = int(usage[proj_name]['quota_limit'])
+                    old_used_space_perc = usage[proj_name]['used_space_perc']
+                if (len((str(limit)).strip()) > 0 and limit > 0):
+                    logger.info("Calculating used space percentage")
+                    project_usage['used_space_perc'] = quota*100/limit
+                else:
+                    logger.info("No storage space allocated")
+                    continue
+                if self.conf.notification_active:
+                    self._thresholdAlarm(project_usage, proj_name, 
+                                         old_used_space_perc)
+                if usage: 
+                    usage[proj_name] = project_usage
+                    total_used_space += quota
+                    total_allocated_space += limit
+                    total_number_of_objects += num_objs
 
         if usage: 
+            logger.info("Writing storage stats to file")
             self._writeUsageStats(usage)
+            logger.info("Total used space: " + str(total_used_space)
+                        + " " + self.conf.storage_space_unity)
+            logger.info("Total allocated space: " + str(total_allocated_space)
+                        + " " + self.conf.storage_space_unity)
+            logger.info("Total number of objects: " + str(total_number_of_objects))
             return True
         else:
+            logger.info("No storage stats to write")
             return False
 
 
@@ -211,18 +293,26 @@ def parseConf(self):
             return False
 
         self.stat_output_file = tmp['stat_output_file']
+        self.storage_space_unity = tmp['storage_space_unity']
         self.internal_project_list = (tmp['internal_project_list']).split(',')
         self.mirrored_projects = (tmp['mirrored_projects']).split(',')
         self.irods_home_dir = tmp['irods_home_dir']
-        self.irods_debug = tmp['irods_debug']
+        if tmp['irods_debug'] in ['True', 'true']:
+             self.irods_debug = True
+        else: 
+            self.irods_debug = False
 
         logger.setLevel(self.log_level[tmp['log_level']])
         rfh = logging.handlers.RotatingFileHandler(tmp['log_file'],maxBytes=8388608,backupCount=9)
         formatter = logging.Formatter('%(asctime)s %(levelname)s: ' 
                                       + '[%(funcName)s] %(message)s')
         rfh.setFormatter(formatter)
         logger.addHandler(rfh)
-
+
+        if tmp['notification_active'] in ['True', 'true']:
+            self.notification_active = True
+        else:
+            self.notification_active = False
         self.notification_sender = tmp['notification_sender']
         self.notification_receiver = tmp['notification_receiver']
 
@@ -246,8 +336,8 @@ def compute(args):
 
 if __name__ == "__main__":
 
-    parser = argparse.ArgumentParser(description='CINECA iRODS quota computation tool')
-    parser.add_argument("confpath",default="NULL",help="path to the configuration file")
+    parser = argparse.ArgumentParser(description='iRODS collection quota computation tool')
+    parser.add_argument("confpath",help="path to the configuration file")
     parser.add_argument("-d", "--dryrun", action="store_true",\
         help="execute a command without performing any real change")