Register resources in quota holder
authorSofia Papagiannaki <papagian@gmail.com>
Mon, 17 Sep 2012 15:36:50 +0000 (18:36 +0300)
committerSofia Papagiannaki <papagian@gmail.com>
Mon, 17 Sep 2012 15:36:50 +0000 (18:36 +0300)
snf-astakos-app/astakos/im/endpoints/quotaholder.py
snf-astakos-app/astakos/im/management/commands/quotaholder_bootstrap.py
snf-astakos-app/astakos/im/migrations/0017_populate_resource_data.py
snf-astakos-app/astakos/im/models.py

index f9f8f58..dc518fc 100644 (file)
 
 import socket
 import logging
+import itertools
+
+from functools import wraps
+from itertools import tee
 
 from django.utils.translation import ugettext as _
 
@@ -46,50 +50,72 @@ ENTITY_KEY = '1'
 logger = logging.getLogger(__name__)
 
 
-def register_users(users, client=None):
-    if not users:
-        return
-
-    if not QUOTA_HOLDER_URL:
-        return
-
-    c = client or QuotaholderHTTP(QUOTA_HOLDER_URL)
+def call(func_name):
+    """Decorator function for QuotaholderHTTP client calls."""
+    def decorator(payload_func):
+        @wraps(payload_func)
+        def wrapper(entities=(), client=None, **kwargs):
+            if not entities:
+                return ()
+        
+            if not QUOTA_HOLDER_URL:
+                return ()
+        
+            c = client or QuotaholderHTTP(QUOTA_HOLDER_URL)
+            func = c.__dict__.get(func_name)
+            if not func:
+                return c,
+            
+            data = payload_func(entities, client, **kwargs)
+            if not data:
+                return c,
+            
+            funcname = func.__name__
+            kwargs = {'context': {}, funcname: data}
+            rejected = func(**kwargs)
+            msg = _('%s: %s - Rejected: %s' % (funcname, data, rejected,))
+            logger.log(LOGGING_LEVEL, msg)
+            return c, rejected
+        return wrapper
+    return decorator
+
+@call('set_quota')
+def send_quota(users, client=None):
     data = []
     append = data.append
     for user in users:
-        try:
-            entity = user.email
-        except AttributeError:
-            continue
-        else:
-            args = entity, owner, key, ownerkey = (
-                entity, 'system', ENTITY_KEY, ''
-            )
+        for resource, limit in user.quota.iteritems():
+            key = ENTITY_KEY
+            quantity = None
+            capacity = limit
+            import_limit = None
+            export_limit = None
+            flags = 0
+            args = (user.email, resource, key, quantity, capacity, import_limit,
+                    export_limit, flags)
             append(args)
+    return data
 
-    if not data:
-        return
-
-    rejected = c.create_entity(
-        context={},
-        create_entity=data,
-    )
-    msg = _('Create entities: %s - Rejected: %s' % (data, rejected,))
-    logger.log(LOGGING_LEVEL, msg)
-
-    created = filter(lambda u: unicode(u.email) not in rejected, users)
-    send_quota(created, c)
-    return rejected
-
-
-def send_quota(users, client=None):
-    if not users:
-        return
-
-    if not QUOTA_HOLDER_URL:
-        return
 
-    c = client or QuotaholderHTTP(QUOTA_HOLDER_URL)
+@call('set_quota')
+def send_resource_quantities(resources, client=None):
+    data = []
+    append = data.append
+    for resource in resources:
+        key = ENTITY_KEY
+        quantity = resource.meta.filter(key='quantity') or None
+        capacity = None
+        import_limit = None
+        export_limit = None
+        flags = 0
+        args = (resource.service, resource, key, quantity, capacity,
+                import_limit, export_limit, flags)
+        append(args)
+    return data
+
+
+@call('get_quota')
+def get_quota(users, client=None):
     data = []
     append = data.append
     for user in users:
@@ -98,46 +124,42 @@ def send_quota(users, client=None):
         except AttributeError:
             continue
         else:
-            for resource, limit in user.quota.iteritems():
-                args = entity, resource, key, quantity, capacity, import_limit, \
-                    export_limit, flags = (
-                        entity, resource, ENTITY_KEY, '0', str(limit), 0, 0, 0
-                    )
+            for r in user.quota.keys():
+                args = entity, r, ENTITY_KEY
                 append(args)
+    return data
 
-    if not data:
-        return
-
-    rejected = c.set_quota(context={}, set_quota=data)
-    msg = _('Set quota: %s - Rejected: %s' % (data, rejected,))
-    logger.log(LOGGING_LEVEL, msg)
-    return rejected
-
-
-def get_quota(users, client=None):
-    if not users:
-        return
 
-    if not QUOTA_HOLDER_URL:
-        return
-
-    c = client or QuotaholderHTTP(QUOTA_HOLDER_URL)
+@call('create_entity')
+def create_entities(entities, client=None, field=''):
     data = []
     append = data.append
-    for user in users:
+    for entity in entities:
         try:
-            entity = user.email
+            entity = entity.__getattribute__(field)
         except AttributeError:
             continue
-        else:
-            for r in user.quota.keys():
-                args = entity, resource, key = entity, r, ENTITY_KEY
-                append(args)
+        owner = 'system'
+        key = ENTITY_KEY
+        ownerkey = ''
+        args = entity, owner, key, ownerkey
+        append(args)
+    return data
 
-    if not data:
-        return
 
-    r = c.get_quota(context={}, get_quota=data)
-    msg = _('Get quota: %s' % data)
-    logger.log(LOGGING_LEVEL, msg)
-    return r
+def register_users(users, client=None):
+    users, copy = itertools.tee(users)
+    client, rejected = create_entities(entities=users,
+                                       client=client,
+                                       field='email')
+    created = (e for e in copy if unicode(e) not in rejected)
+    return send_quota(created, client)
+
+
+def register_resources(resources, client=None):
+    resources, copy = itertools.tee(resources)
+    client, rejected = create_entities(entities=resources,
+                                       client=client,
+                                       field='service')
+    created = (e for e in copy if unicode(e) not in rejected)
+    return send_resource_quantities(created, client)
\ No newline at end of file
index ef48f12..a24a8ab 100644 (file)
@@ -34,8 +34,8 @@
 from django.core.management.base import BaseCommand, CommandError
 from django.db.utils import IntegrityError
 
-from astakos.im.models import AstakosUser
-from astakos.im.endpoints.quotaholder import register_users
+from astakos.im.models import AstakosUser, Resource
+from astakos.im.endpoints.quotaholder import register_users, register_resources
 
 
 class Command(BaseCommand):
@@ -43,8 +43,7 @@ class Command(BaseCommand):
 
     def handle(self, *args, **options):
         try:
-            r = register_users(AstakosUser.objects.all())
-            self.stdout.write("Rejected: %s\n" % r)
+            register_resources(Resource.objects.all())
+            register_users(AstakosUser.objects.all())
         except BaseException, e:
-            print e
             raise CommandError("Bootstrap failed.")
index 8be9327..e69cbcb 100644 (file)
@@ -4,7 +4,7 @@ from south.v2 import DataMigration
 
 d = {
     'cyclades': {'vm': 2},
-    'pithos+': {'diskspace': 10}
+    'pithos+': {'diskspace': 50 * 1024 * 1024 * 1024} # 5 GB
 }
 
 
index a628233..1b72608 100644 (file)
@@ -51,9 +51,10 @@ from django.dispatch import Signal
 from django.db.models import Q
 
 from astakos.im.settings import (DEFAULT_USER_LEVEL, INVITATIONS_PER_LEVEL,
-                                 AUTH_TOKEN_DURATION, BILLING_FIELDS, EMAILCHANGE_ACTIVATION_DAYS, LOGGING_LEVEL
-                                 )
-from astakos.im.endpoints.quotaholder import register_users, send_quota
+                                 AUTH_TOKEN_DURATION, BILLING_FIELDS,
+                                 EMAILCHANGE_ACTIVATION_DAYS, LOGGING_LEVEL)
+from astakos.im.endpoints.quotaholder import (register_users, send_quota,
+                                              register_resources)
 from astakos.im.endpoints.aquarium.producer import report_user_event
 
 from astakos.im.tasks import propagate_groupmembers_quota
@@ -616,6 +617,12 @@ def astakosuser_post_save(sender, instance, created, **kwargs):
     register_users((instance,))
 
 
+def resource_post_save(sender, instance, created, **kwargs):
+    if not created:
+        return
+    register_resources((instance,))
+
+
 def send_quota_disturbed(sender, instance, **kwargs):
     users = []
     extend = users.extend
@@ -645,6 +652,7 @@ post_syncdb.connect(fix_superusers)
 post_save.connect(user_post_save, sender=User)
 pre_save.connect(astakosuser_pre_save, sender=AstakosUser)
 post_save.connect(astakosuser_post_save, sender=AstakosUser)
+post_save.connect(resource_post_save, sender=Resource)
 
 quota_disturbed = Signal(providing_args=["users"])
 quota_disturbed.connect(on_quota_disturbed)