Revision 45112d5a snf-astakos-app/astakos/im/management/commands/user-modify.py

b/snf-astakos-app/astakos/im/management/commands/user-modify.py
43 43
from django.core.validators import validate_email
44 44

  
45 45
from synnefo.util import units
46
from astakos.im.models import AstakosUser, AstakosUserQuota
46
from astakos.im.models import AstakosUser, Resource
47 47
from astakos.im import quotas
48 48
from astakos.im import activation_backends
49 49
from ._common import (remove_user_permission, add_user_permission, is_uuid,
......
53 53

  
54 54

  
55 55
class Command(BaseCommand):
56
    args = "<user ID>"
56
    args = "<user ID> (or --all)"
57 57
    help = "Modify a user's attributes"
58 58

  
59 59
    option_list = BaseCommand.option_list + (
60
        make_option('--all',
61
                    action='store_true',
62
                    default=False,
63
                    help=("Operate on all users. Currently only setting "
64
                          "base quota is supported in this mode. Can be "
65
                          "combined with `--exclude'.")),
66
        make_option('--exclude',
67
                    help=("If `--all' is given, exclude users given as a "
68
                          "list of uuids: uuid1,uuid2,uuid3")),
60 69
        make_option('--invitations',
61 70
                    dest='invitations',
62 71
                    metavar='NUM',
......
129 138
        make_option('--reject-reason',
130 139
                    dest='reject_reason',
131 140
                    help="Reason user got rejected"),
132
        make_option('--set-base-quota',
141
        make_option('--base-quota',
133 142
                    dest='set_base_quota',
134 143
                    metavar='<resource> <capacity>',
135 144
                    nargs=2,
......
153 162

  
154 163
    @transaction.commit_on_success
155 164
    def handle(self, *args, **options):
165
        if options['all']:
166
            if not args:
167
                return self.handle_all_users(*args, **options)
168
            else:
169
                raise CommandError("Please provide a user ID or --all")
170

  
156 171
        if len(args) != 1:
157
            raise CommandError("Please provide a user ID")
172
            raise CommandError("Please provide a user ID or --all")
173

  
174
        if options["exclude"] is not None:
175
            m = "Option --exclude is meaningful only combined with --all."
176
            raise CommandError(m)
158 177

  
159 178
        if args[0].isdigit():
160 179
            try:
161
                user = AstakosUser.objects.get(id=int(args[0]))
180
                user = AstakosUser.objects.select_for_update().\
181
                    get(id=int(args[0]))
162 182
            except AstakosUser.DoesNotExist:
163 183
                raise CommandError("Invalid user ID")
164 184
        elif is_uuid(args[0]):
......
304 324
        set_base_quota = options.get('set_base_quota')
305 325
        if set_base_quota is not None:
306 326
            resource, capacity = set_base_quota
307
            self.set_limit(user, resource, capacity, force)
327
            self.set_limits([user], resource, capacity, force)
308 328

  
309 329
        delete = options.get('delete')
310 330
        if delete:
......
335 355
            user.email = newemail
336 356
            user.save()
337 357

  
338
    def set_limit(self, user, resource, capacity, force):
358
    def confirm(self):
359
        self.stdout.write("Confirm? [y/N] ")
360
        response = raw_input()
361
        if string.lower(response) not in ['y', 'yes']:
362
            self.stdout.write("Aborted.\n")
363
            exit()
364

  
365
    def handle_limits_user(self, user, res, capacity, style):
366
        default_capacity = res.uplimit
367
        resource = res.name
368
        quota = user.get_resource_policy(resource)
369
        s_default = show_resource_value(default_capacity, resource, style)
370
        s_current = show_resource_value(quota.capacity, resource, style)
371
        s_capacity = (show_resource_value(capacity, resource, style)
372
                      if capacity != 'default' else capacity)
373
        self.stdout.write("user: %s (%s)\n" % (user.uuid, user.username))
374
        self.stdout.write("default capacity: %s\n" % s_default)
375
        self.stdout.write("current capacity: %s\n" % s_current)
376
        self.stdout.write("new capacity: %s\n" % s_capacity)
377
        self.confirm()
378

  
379
    def handle_limits_all(self, res, capacity, exclude, style):
380
        m = "This will set base quota for all users"
381
        app = (" except %s" % ", ".join(exclude)) if exclude else ""
382
        self.stdout.write(m+app+".\n")
383
        resource = res.name
384
        self.stdout.write("resource: %s\n" % resource)
385
        s_capacity = (show_resource_value(capacity, resource, style)
386
                      if capacity != 'default' else capacity)
387
        self.stdout.write("capacity: %s\n" % s_capacity)
388
        self.confirm()
389

  
390
    def set_limits(self, users, resource, capacity, force=False, exclude=None):
391
        try:
392
            r = Resource.objects.get(name=resource)
393
        except Resource.DoesNotExist:
394
            raise CommandError("No such resource '%s'." % resource)
395

  
339 396
        style = None
340
        if capacity != 'default':
397
        if capacity != "default":
341 398
            try:
342 399
                capacity, style = units.parse_with_style(capacity)
343 400
            except:
344
                m = "Please specify capacity as a decimal integer or 'default'"
401
                m = ("Please specify capacity as a decimal integer or "
402
                     "'default'")
345 403
                raise CommandError(m)
346 404

  
347
        try:
348
            quota = user.get_resource_policy(resource)
349
        except AstakosUserQuota.DoesNotExist:
350
            raise CommandError("No such resource: %s" % resource)
351

  
352
        default_capacity = quota.resource.uplimit
353 405
        if not force:
354
            s_default = show_resource_value(default_capacity, resource, style)
355
            s_current = show_resource_value(quota.capacity, resource, style)
356
            s_capacity = (show_resource_value(capacity, resource, style)
357
                          if capacity != 'default' else capacity)
358
            self.stdout.write("user: %s (%s)\n" % (user.uuid, user.username))
359
            self.stdout.write("default capacity: %s\n" % s_default)
360
            self.stdout.write("current capacity: %s\n" % s_current)
361
            self.stdout.write("new capacity: %s\n" % s_capacity)
362
            self.stdout.write("Confirm? (y/n) ")
363
            response = raw_input()
364
            if string.lower(response) not in ['y', 'yes']:
365
                self.stdout.write("Aborted.\n")
366
                return
367

  
368
        if capacity == 'default':
369
            capacity = default_capacity
370
        quotas.update_base_quota(quota, capacity)
406
            if len(users) == 1:
407
                self.handle_limits_user(users[0], r, capacity, style)
408
            else:
409
                self.handle_limits_all(r, capacity, exclude, style)
410

  
411
        if capacity == "default":
412
            capacity = r.uplimit
413
        quotas.update_base_quota(users, resource, capacity)
414

  
415
    def handle_all_users(self, *args, **options):
416
        force = options["force"]
417
        exclude = options["exclude"]
418
        if exclude is not None:
419
            exclude = exclude.split(',')
420

  
421
        set_base_quota = options.get('set_base_quota')
422
        if set_base_quota is not None:
423
            users = AstakosUser.objects.accepted().select_for_update()
424
            if exclude:
425
                users = users.exclude(uuid__in=exclude)
426
            resource, capacity = set_base_quota
427
            self.set_limits(users, resource, capacity, force, exclude)

Also available in: Unified diff