Revision 3a527b3a snf-astakos-app/astakos/im/management/commands/resource-modify.py

b/snf-astakos-app/astakos/im/management/commands/resource-modify.py
33 33

  
34 34
from optparse import make_option
35 35
from django.core.management.base import BaseCommand, CommandError
36
from django.utils import simplejson as json
36 37

  
37 38
from astakos.im.models import Resource
38 39
from astakos.im.resources import update_resource
......
40 41

  
41 42
class Command(BaseCommand):
42 43
    args = "<resource name>"
43
    help = "Modify a resource (currently only change the default base quota)"
44
    help = ("Modify a resource (currently only change the default base quota)."
45
            "\nIf no resource is specified, all resources are considered.")
44 46

  
45 47
    option_list = BaseCommand.option_list + (
46 48
        make_option('--limit',
47 49
                    dest='limit',
48
                    help="Change default base quota"),
50
                    help="Specify default base quota"),
51
        make_option('--interactive',
52
                    action='store_true',
53
                    dest='interactive',
54
                    default=None,
55
                    help="Prompt user to change default base quotas"),
56
        make_option('--from-file',
57
                    dest='from_file',
58
                    metavar='<limits_file.json>',
59
                    help="Read default base quotas from a json file"),
49 60
    )
50 61

  
51 62
    def handle(self, *args, **options):
52
        if len(args) < 1:
63
        resource_name = args[0] if len(args) > 0 else None
64

  
65
        actions = {
66
            'limit': self.change_limit,
67
            'interactive': self.change_interactive,
68
            'from_file': self.change_from_file,
69
        }
70

  
71
        opts = [(key, value)
72
                for (key, value) in options.items()
73
                if key in actions and value is not None]
74

  
75
        if len(opts) != 1:
76
            raise CommandError("Please provide exactly one option.")
77

  
78
        key, value = opts[0]
79
        action = actions[key]
80
        action(resource_name, value)
81

  
82
    def get_resource(self, resource_name):
83
        try:
84
            return Resource.objects.get_for_update(name=resource_name)
85
        except Resource.DoesNotExist:
86
            raise CommandError("Resource %s does not exist."
87
                               % resource_name)
88

  
89
    def change_limit(self, resource_name, limit):
90
        if resource_name is None:
53 91
            raise CommandError("Please provide a resource name.")
54 92

  
55
        resource_name = args[0]
93
        resource = self.get_resource(resource_name)
94
        self.change_resource_limit(resource, limit)
56 95

  
57
        limit = options['limit']
58
        if limit is None:
59
            raise CommandError("Use --limit to change default base quota.")
96
    def change_from_file(self, resource_name, filename):
97
        with open(filename) as file_data:
98
            try:
99
                config = json.load(file_data)
100
            except json.JSONDecodeError:
101
                raise CommandError("Malformed JSON file.")
102
            if not isinstance(config, dict):
103
                raise CommandError("Malformed JSON file.")
104
            self.change_with_conf(resource_name, config)
60 105

  
106
    def change_with_conf(self, resource_name, config):
107
        if resource_name is None:
108
            resources = Resource.objects.all().select_for_update()
109
        else:
110
            resources = [self.get_resource(resource_name)]
111

  
112
        for resource in resources:
113
            limit = config.get(resource.name)
114
            if limit is not None:
115
                self.change_resource_limit(resource, limit)
116

  
117
    def change_interactive(self, resource_name, _placeholder):
118
        if resource_name is None:
119
            resources = Resource.objects.all().select_for_update()
120
        else:
121
            resources = [self.get_resource(resource_name)]
122

  
123
        for resource in resources:
124
            self.stdout.write("Resource '%s' (%s)\n" %
125
                              (resource.name, resource.desc))
126
            unit = (" in %s" % resource.unit) if resource.unit else ""
127
            self.stdout.write("Current limit%s: %s\n"
128
                              % (unit, resource.uplimit))
129
            while True:
130
                self.stdout.write("New limit%s (leave blank to keep current): "
131
                                  % (unit))
132
                response = raw_input()
133
                if response == "":
134
                    break
135
                else:
136
                    try:
137
                        value = int(response)
138
                    except ValueError:
139
                        continue
140
                    update_resource(resource, value)
141
                    break
142

  
143
    def change_resource_limit(self, resource, limit):
61 144
        try:
62 145
            limit = int(limit)
63 146
        except:
64 147
            raise CommandError("Limit should be an integer.")
65

  
66
        try:
67
            update_resource(resource_name, limit)
68
        except Resource.DoesNotExist:
69
            raise CommandError("Resource %s does not exist." % resource_name)
148
        update_resource(resource, limit)

Also available in: Unified diff