X-Git-Url: http://xvm.mit.edu/gitweb/invirt/packages/invirt-base.git/blobdiff_plain/bc9a52eeb172e8e416b3d8c576a722ca485a6652..9014d0616c6c4fe5f2e19b2a3889700cf70dd654:/scripts/invirt-setquotas

diff --git a/scripts/invirt-setquotas b/scripts/invirt-setquotas
index af33eec..ceae1bf 100755
--- a/scripts/invirt-setquotas
+++ b/scripts/invirt-setquotas
@@ -7,110 +7,100 @@ that owner.  Setting a parameter to -1 restores the default.
 
 Examples:
 
-    invirt-setquota joeuser -mt 512 -ms None
+    invirt-setquota joeuser -mt 512 -ms -1
 """
 
 from invirt.database import *
 from sys import argv, exit, stderr, stdout
 from optparse import OptionParser
 
-class invirt_exception(Exception): pass
-
 def main(argv):
-    try:
-        parser = OptionParser(usage = '%prog owner [options]',
-                description = __doc__.strip().split('\n\n')[0])
-        parser.add_option('-m', '--mem-total',
-                type = 'int',
-                dest = 'memtotal',
-                help = 'set total concurrent RAM quota')
-        parser.add_option('-n', '--mem-single',
-                type = 'int',
-                dest = 'memsingle',
-                help = 'set single VM RAM quota')
-        parser.add_option('-d', '--disk-total',
-                type = 'int',
-                dest = 'disktotal',
-                help = 'set total disk quota')
-        parser.add_option('-e', '--disk-single',
-                type = 'int',
-                dest = 'disksingle',
-                help = 'set single VM disk quota')
-        parser.add_option('-v', '--vms-total',
-                type = 'int',
-                dest = 'vmstotal',
-                help = 'set total VM quota')
-        parser.add_option('-w', '--vms-active',
-                type = 'int',
-                dest = 'vmsactive',
-                help = 'set active VM quota')
-        opts, args = parser.parse_args()
-
-        print opts
-        if len(args) != 1:
-            raise invirt_exception(__doc__.strip())
-        owner = args[0]
-        connect()
-        session.begin()
-        
-        x = Owner.query().filter_by(owner_id=owner).first()
-
-        edited = False
-        if opts.memtotal != None:
-            total = int(opts.memtotal)
-            if total == -1:
-                x.ram_quota_total = None
-            else:
-                x.ram_quota_total = total
-            edited = True
-
-        if opts.memsingle != None:
-            single = int(opts.memsingle)
-            if single == -1:
-                x.ram_quota_single = None
-            else:
-                x.ram_quota_single = single
-            edited = True
-
-        if opts.disktotal != None:
-            total = int(opts.disktotal)
-            if total == -1:
-                x.disk_quota_total = None
-            else:
-                x.disk_quota_total = total
-            edited = True
-
-        if opts.disksingle != None:
-            single = int(opts.disksingle)
-            if single == -1:
-                x.disk_quota_single = None
-            else:
-                x.disk_quota_single = single
-            edited = True
-
-        if opts.vmstotal != None:
-            total = int(opts.vmstotal)
-            if total == -1:
-                x.vms_quota_total = None
-            else:
-                x.vms_quota_total = total
-            edited = True
-
-        if opts.vmsactive != None:
-            active = int(opts.vmsactive)
-            if active == -1:
-                x.vms_quota_active = None
-            else:
-                x.vms_quota_active = active
-            edited = True
-
-        if edited:
-            session.commit()
-        print str(x)
-
-    except invirt_exception, ex:
-        print >> stderr, ex
+    parser = OptionParser(usage = '%prog owner [options]',
+            description = __doc__.strip().split('\n\n')[0])
+    parser.add_option('-m', '--mem-total',
+            type = 'int',
+            dest = 'memtotal',
+            help = 'set total concurrent RAM quota')
+    parser.add_option('-n', '--mem-single',
+            type = 'int',
+            dest = 'memsingle',
+            help = 'set single VM RAM quota')
+    parser.add_option('-d', '--disk-total',
+            type = 'int',
+            dest = 'disktotal',
+            help = 'set total disk quota')
+    parser.add_option('-e', '--disk-single',
+            type = 'int',
+            dest = 'disksingle',
+            help = 'set single VM disk quota')
+    parser.add_option('-v', '--vms-total',
+            type = 'int',
+            dest = 'vmstotal',
+            help = 'set total VM quota')
+    parser.add_option('-w', '--vms-active',
+            type = 'int',
+            dest = 'vmsactive',
+            help = 'set active VM quota')
+    opts, args = parser.parse_args()
+
+    if len(args) != 1:
+        print >> stderr, __doc__.strip()
         return 1
+    owner = args[0]
+    connect()
+    session.begin()
+    
+    x = Owner.query().filter_by(owner_id=owner).first()
+    if x == None:
+        x = Owner(owner_id=owner, ram_quota_total=None, ram_quota_single=None,
+                  disk_quota_total=None, disk_quota_single=None,
+                  vms_quota_total=None, vms_quota_active=None)
+
+    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
+
+    session.commit()
+    print str(x)
+    return 0
 
 if __name__ == '__main__':
     exit(main(argv))