Revision 9eafaa32 snf-astakos-app/astakos/im/management/commands/user-update.py

b/snf-astakos-app/astakos/im/management/commands/user-update.py
38 38
from django.core.exceptions import ValidationError
39 39
from django.db.utils import IntegrityError
40 40

  
41
from astakos.im.models import AstakosUser, AstakosGroup, Membership
41
from astakos.im.models import (AstakosUser, AstakosGroup, Membership, Resource,
42
                               AstakosUserQuota)
42 43
from astakos.im.endpoints.aquarium.producer import report_user_credits_event
43 44
from ._common import remove_user_permission, add_user_permission
44 45

  
......
47 48
    args = "<user ID>"
48 49
    help = "Modify a user's attributes"
49 50

  
50
    option_list = BaseCommand.option_list + (
51
    option_list = list(BaseCommand.option_list) + [
51 52
        make_option('--invitations',
52 53
                    dest='invitations',
53 54
                    metavar='NUM',
......
111 112
                    dest='refill',
112 113
                    default=False,
113 114
                    help="Refill user credits"),
114
    )
115

  
115
    ]
116
    resources = Resource.objects.select_related().all()
117
    append = option_list.append
118
    for r in resources:
119
        append(make_option('--%s-set-quota' % r,
120
                    dest='%s-set-quota' % r,
121
                    metavar='QUANTITY',
122
                    help="Set resource quota"))
123
    
116 124
    def handle(self, *args, **options):
117 125
        if len(args) != 1:
118 126
            raise CommandError("Please provide a user ID")
......
225 233

  
226 234
        if password:
227 235
            self.stdout.write('User\'s new password: %s\n' % password)
236
        
237
        for r in self.resources:
238
            limit = options.get('%s-set-quota' % r)
239
            if not limit:
240
                continue
241
            if not limit.isdigit():
242
                raise CommandError('Invalid limit')
243
            
244
            q = AstakosUserQuota.objects
245
            q, created = q.get_or_create(resource=r, user=user,
246
                                         defaults={'uplimit': limit})
247
            verb = 'set' if created else 'updated'
248
            self.stdout.write('User\'s quota %s successfully\n' % verb)

Also available in: Unified diff