source: trunk/packages/invirt-database/scripts/invirt-quota @ 2162

#!/usr/bin/env python

"""
invirt-setquota 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-setquota joeuser -mt 512 -ms -1
"""

from invirt.database import *
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',
            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)

    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))