invirt-quota: refactor; print full help
[invirt/packages/invirt-database.git] / scripts / invirt-quota
index 73fef05..8a7f67c 100755 (executable)
@@ -4,10 +4,6 @@
 invirt-quota allows an administrator to set memory, disk, and VM quotas 
 for an owner.  Invoking with only an owner name returns the current quotas for
 that owner.  Setting a parameter to -1 restores the default.
-
-Examples:
-
-    invirt-quota joeuser -mt 512 -ms -1
 """
 
 from invirt.database import *
@@ -15,15 +11,15 @@ from sys import argv, exit, stderr, stdout
 from optparse import OptionParser
 
 def main(argv):
-    parser = OptionParser(usage = '%prog owner [options]',
-            description = __doc__.strip().split('\n\n')[0])
-    parser.add_option('-m', '--mem-total',
+    parser = OptionParser(usage = '%prog <owner> [options]',
+            description = __doc__.strip())
+    parser.add_option('-m', '--ram-total',
             type = 'int',
-            dest = 'memtotal',
+            dest = 'ramtotal',
             help = 'set total concurrent RAM quota')
-    parser.add_option('-n', '--mem-single',
+    parser.add_option('-n', '--ram-single',
             type = 'int',
-            dest = 'memsingle',
+            dest = 'ramsingle',
             help = 'set single VM RAM quota')
     parser.add_option('-d', '--disk-total',
             type = 'int',
@@ -44,57 +40,23 @@ def main(argv):
     opts, args = parser.parse_args()
 
     if len(args) != 1:
-        print >> stderr, __doc__.strip()
+        parser.print_help(stderr)
         return 1
     owner = args[0]
     connect()
     session.begin()
     
     x = Owner.query().filter_by(owner_id=owner).first()
-    if x == None:
+    if x is None:
         x = Owner(owner_id=owner)
 
-    if opts.memtotal != None:
-        total = int(opts.memtotal)
-        if total == -1:
-            x.ram_quota_total = None
-        else:
-            x.ram_quota_total = total
-
-    if opts.memsingle != None:
-        single = int(opts.memsingle)
-        if single == -1:
-            x.ram_quota_single = None
-        else:
-            x.ram_quota_single = single
-
-    if opts.disktotal != None:
-        total = int(opts.disktotal)
-        if total == -1:
-            x.disk_quota_total = None
-        else:
-            x.disk_quota_total = total
-
-    if opts.disksingle != None:
-        single = int(opts.disksingle)
-        if single == -1:
-            x.disk_quota_single = None
-        else:
-            x.disk_quota_single = single
-
-    if opts.vmstotal != None:
-        total = int(opts.vmstotal)
-        if total == -1:
-            x.vms_quota_total = None
-        else:
-            x.vms_quota_total = total
-
-    if opts.vmsactive != None:
-        active = int(opts.vmsactive)
-        if active == -1:
-            x.vms_quota_active = None
-        else:
-            x.vms_quota_active = active
+    for resource, scope in [('ram',  'total'), ('ram',  'single'),
+                            ('disk', 'total'), ('disk', 'single'),
+                            ('vms',  'total'), ('vms',  'active')]:
+        opt = getattr(opts, resource+scope)
+        if opt is not None:
+            val = int(opt)
+            setattr(x, resource+'_quota_'+scope, val if val >= 0 else None)
 
     session.commit()
     print str(x)