[#474][#479] Allow querying, setting, and removing quotas

Author: d-w-moore

irods/manager/user_manager.py

@@ -5,7 +5,8 @@
 
 from irods.models import User, Group
 from irods.manager import Manager
-from irods.message import UserAdminRequest, GeneralAdminRequest, iRODSMessage, GetTempPasswordForOtherRequest, GetTempPasswordForOtherOut
+from irods.message import (UserAdminRequest, GeneralAdminRequest, _do_GeneralAdminRequest, iRODSMessage,
+        GetTempPasswordForOtherRequest, GetTempPasswordForOtherOut)
 from irods.exception import UserDoesNotExist, GroupDoesNotExist, NoResultFound, CAT_SQL_ERR
 from irods.api_number import api_number
 from irods.user import iRODSUser, iRODSGroup, Bad_password_change_parameter
@@ -16,6 +17,30 @@
 
 class UserManager(Manager):
 
+    def _get_session(self): return self.sess
+
+    def calculate_usage(self):
+        return _do_GeneralAdminRequest(self._get_session, "calculate-usage")
+
+    # TODO: remove this in branch 2.x (#482)
+    def set_quota(self, user_name, amount, resource = 'total'):
+        return _do_GeneralAdminRequest(self._get_session,
+            "set-quota",
+            "user",
+            user_name,
+            resource,
+            str(amount)
+        )
+
+    def remove_quota(self, user_name, resource = 'total'):
+        return _do_GeneralAdminRequest(self._get_session,
+            "set-quota",
+            "user",
+            user_name,
+            resource,
+            "0"
+        )
+
     def get(self, user_name, user_zone=""):
         query = self.sess.query(User).filter(User.name == user_name)
 
@@ -319,6 +344,24 @@ def removemember(self,
             response = conn.recv()
         logger.debug(response.int_info)
 
+    def remove_quota(self, group_name, resource = 'total'):
+        self.set_quota(group_name, amount = 0, resource = resource)
+
+    def set_quota(self, group_name, amount, resource = 'total'):
+        message_body = GeneralAdminRequest(
+            "set-quota",
+            "group",
+            group_name,
+            resource,
+            str(amount)
+        )
+        request = iRODSMessage("RODS_API_REQ", msg=message_body,
+                               int_info=api_number['GENERAL_ADMIN_AN'])
+        with self.sess.pool.get_connection() as conn:
+            conn.send(request)
+            response = conn.recv()
+        logger.debug(response.int_info)
+
 
 # NOTE: Everything of the form *UserGroup* is deprecated.
 UserGroupManager = GroupManager

irods/message/__init__.py

@@ -10,6 +10,7 @@
 import xml.etree.ElementTree as ET_xml
 import defusedxml.ElementTree as ET_secure_xml
 from . import quasixml as ET_quasi_xml
+from ..api_number import api_number
 from collections import namedtuple
 import os
 import ast
@@ -784,6 +785,18 @@ def __init__(self, *args):
 class GeneralAdminRequest(_admin_request_base):
     _name = 'generalAdminInp_PI'
 
+def _do_GeneralAdminRequest(_session_or_accessor, *args):
+    sess = _session_or_accessor
+    if callable(sess):
+        sess = sess()
+    message_body = GeneralAdminRequest(*args)
+    request = iRODSMessage("RODS_API_REQ", msg=message_body,
+                           int_info=api_number['GENERAL_ADMIN_AN'])
+    with sess.pool.get_connection() as conn:
+        conn.send(request)
+        response = conn.recv()
+    logger.debug(response.int_info)
+
 
 # define userAdminInp_PI "str *arg0; str *arg1; str *arg2; str *arg3;
 # str *arg4; str *arg5; str *arg6; str *arg7;  str *arg8;  str *arg9;"

irods/models.py

@@ -199,6 +199,14 @@ class Keywords(Model):
     chksum = Keyword(String, 'chksum')
 
 
+class Quota(Model):
+    user_id = Column(Integer, 'COL_QUOTA_USER_ID', 2000)
+    resc_id = Column(Integer, 'COL_QUOTA_RESC_ID', 2001)
+    limit = Column(Integer, 'COL_QUOTA_LIMIT', 2002)
+    over = Column(Integer, 'COL_QUOTA_OVER', 2003)
+    modify_time = Column(DateTime,'COL_QUOTA_MODIFY_TIME', 2004)
+
+
 class TicketQuery:
     """Various model classes for querying attributes of iRODS tickets.
 

irods/test/user_group_test.py

@@ -6,17 +6,43 @@
 import unittest
 import tempfile
 import shutil
+from subprocess import check_output, CalledProcessError
 from irods import MAX_PASSWORD_LENGTH
 from irods.exception import GroupDoesNotExist, UserDoesNotExist
 from irods.meta import iRODSMetaCollection, iRODSMeta
-from irods.models import User, Group, UserMeta
+from irods.models import User, Group, UserMeta, Quota
 from irods.session import iRODSSession
 import irods.exception as ex
 import irods.test.helpers as helpers
-from irods.user import Bad_password_change_parameter
+from irods.user import iRODSUser, Bad_password_change_parameter
 from six.moves import range
 
-class TestGroup(unittest.TestCase):
+def user_quotas_implemented():
+    return getattr(iRODSUser, "set_quota", None) is not None
+
+def icommands_available():
+    try:
+        which_output = check_output(['which', 'ihelp'])
+        return b'ihelp' in which_output
+    except CalledProcessError:
+        pass # 'which' returned nonzero status
+    return False
+
+def set_user_quota(session, user_name, bytes_total):
+    success = False
+    if user_quotas_implemented():
+        try:
+            session.users.set_quota(user_name, bytes_total)
+            success = True
+        except:
+            pass
+    if not success and icommands_available():
+        return_code = os.system("iadmin suq {} total {}".format(user_name, bytes_total))
+        if return_code == 0:
+            success = True
+    return success
+
+class TestUserAndGroup(unittest.TestCase):
 
     def setUp(self):
         self.sess = helpers.make_session()
@@ -525,6 +551,60 @@ def alice_login(pw):
             if user_alice:
                 user_alice.remove()
 
+    def test_set_and_query_group_quota(self):
+        #import pdb
+        #pdb.set_trace()
+        bob = quota_group = None
+        ses = self.sess
+        test_object = None
+        my_quota = 10000
+        my_object_size = 1400
+        try:
+            bob = self.sess.users.create('bob', 'rodsuser')
+            ses.users.modify('bob', 'password', 'bpass')
+            quota_group = self.sess.groups.create('quota_group_for_bob')
+            quota_group.addmember(bob.name)
+            quota_group.set_quota(my_quota)
+            with iRODSSession(user = 'bob', password = 'bpass', host = ses.host, port = ses.port, zone = ses.zone) as user_sess:
+                test_object = user_sess.data_objects.create('/tempZone/home/public/bob_file_testing_group_quota')
+                with test_object.open('w') as f:
+                    f.write(b'_' * my_object_size)
+            ses.groups.calculate_usage()
+            for i in ses.query(Quota).filter(Quota.user_id == quota_group.id):
+                self.assertEqual(int(i[Quota.over]), -my_quota + my_object_size)
+        finally:
+            if test_object:
+                test_object.unlink(force = True)
+            if quota_group:
+                quota_group.remove()
+            if bob:
+                ses.users.remove(bob.name)
+
+    @unittest.skipIf(not user_quotas_implemented()
+                     and not icommands_available(), "Can't set up user quota for removal in test")
+    def test_set_and_query_user_quota(self):
+        bob = None
+        if self.sess.server_version >= (4, 3):
+            self.skipTest('iRODS servers 4.3.0 and higher have dropped user quotas in favor of group quotas.')
+        ses = self.sess
+        test_object = None
+        try:
+            bob = self.sess.users.create('bob', 'rodsuser')
+            ses.users.modify('bob', 'password', 'bpass')
+            set_user_quota(ses, 'bob', 10000)
+            with iRODSSession(user = 'bob', password = 'bpass', host = ses.host, port = ses.port, zone = ses.zone) as user_sess:
+               test_object = user_sess.data_objects.create('/tempZone/home/public/bobfile')
+               with test_object.open('w') as f:
+                   f.write(b'_'*1000)
+            ses.users.calculate_usage()
+            for i in ses.query(Quota).filter(Quota.user_id == bob.id):
+                self.assertEqual(int(i[Quota.over]), -9000)
+        finally:
+            if test_object:
+                test_object.unlink(force = True)
+            if bob:
+                ses.users.remove(bob.name)
+
 
 if __name__ == '__main__':
     # let the tests find the parent irods lib

irods/user.py

@@ -9,6 +9,13 @@ class Bad_password_change_parameter(Exception): pass
 
 class iRODSUser(object):
 
+    def remove_quota(self, resource = 'total'):
+        self.manager.remove_quota(self.name, resource = resource)
+
+    # TODO: remove this in branch 2.x (#482)
+    def set_quota(self, amount, resource = 'total'):
+        self.manager.set_quota(self.name, amount, resource = resource)
+
     def __init__(self, manager, result=None):
         self.manager = manager
         if result:
@@ -66,6 +73,12 @@ def temp_password(self):
 
 class iRODSGroup(object):
 
+    def remove_quota(self, resource = 'total'):
+        self.set_quota(amount = 0, resource = resource)
+
+    def set_quota(self, amount, resource = 'total'):
+        self.manager.set_quota(self.name, amount, resource = resource)
+
     def __init__(self, manager, result=None):
         self.manager = manager
         if result:

irods/version.py

@@ -1,4 +1,4 @@
-import os 
+import os
 
 __version__ = '1.1.9'