Revision 764d99c4 snf-astakos-app/astakos/im/management/commands/user-modify.py

b/snf-astakos-app/astakos/im/management/commands/user-modify.py
39 39

  
40 40
from astakos.im.models import AstakosUser
41 41
from astakos.im.functions import (activate, deactivate)
42
from astakos.im import quotas
42 43
from ._common import remove_user_permission, add_user_permission, is_uuid
43 44
from snf_django.lib.db.transaction import commit_on_success_strict
45
import string
44 46

  
45 47

  
46 48
class Command(BaseCommand):
......
102 104
        make_option('--delete-permission',
103 105
                    dest='delete-permission',
104 106
                    help="Delete user permission"),
107
        make_option('--set-base-quota',
108
                    dest='set_base_quota',
109
                    metavar='<resource> <capacity>',
110
                    nargs=2,
111
                    help=("Set base quota for a specified resource. "
112
                          "The special value 'default' sets the user base "
113
                          "quota to the default value.")
114
                    ),
115

  
105 116
    )
106 117

  
107 118
    @commit_on_success_strict()
......
210 221

  
211 222
        if password:
212 223
            self.stdout.write('User\'s new password: %s\n' % password)
224

  
225
        set_base_quota = options.get('set_base_quota')
226
        if set_base_quota is not None:
227
            resource, capacity = set_base_quota
228
            self.set_limit(user, resource, capacity, False)
229

  
230
    def set_limit(self, user, resource, capacity, force):
231
        if capacity != 'default':
232
            try:
233
                capacity = int(capacity)
234
            except ValueError:
235
                m = "Please specify capacity as a decimal integer or 'default'"
236
                raise CommandError(m)
237

  
238
        try:
239
            quota, default_capacity = user.get_resource_policy(resource)
240
        except Resource.DoesNotExist:
241
            raise CommandError("No such resource: %s" % resource)
242

  
243
        current = quota.capacity if quota is not None else 'default'
244

  
245
        if not force:
246
            self.stdout.write("user: %s (%s)\n" % (user.uuid, user.username))
247
            self.stdout.write("default capacity: %s\n" % default_capacity)
248
            self.stdout.write("current capacity: %s\n" % current)
249
            self.stdout.write("new capacity: %s\n" % capacity)
250
            self.stdout.write("Confirm? (y/n) ")
251
            response = raw_input()
252
            if string.lower(response) not in ['y', 'yes']:
253
                self.stdout.write("Aborted.\n")
254
                return
255

  
256
        if capacity == 'default':
257
            try:
258
                quotas.remove_base_quota(user, resource)
259
            except Exception as e:
260
                import traceback
261
                traceback.print_exc()
262
                raise CommandError("Failed to remove policy: %s" % e)
263
        else:
264
            try:
265
                quotas.add_base_quota(user, resource, capacity)
266
            except Exception as e:
267
                raise CommandError("Failed to add policy: %s" % e)

Also available in: Unified diff