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

#!/usr/bin/env python

"""
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.
"""

import sys
import optparse

from invirt import database

def main(argv):
    parser = optparse.OptionParser(usage = '%prog <owner> [options]',
            description = __doc__.strip())
    parser.add_option('-m', '--ram-total',
            type = 'int',
            dest = 'ramtotal',
            help = 'set total concurrent RAM quota')
    parser.add_option('-n', '--ram-single',
            type = 'int',
            dest = 'ramsingle',
            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:
        parser.print_help(sys.stderr)
        return 1
    owner_id = args[0]
    database.connect()
    database.session.begin()
    
    owner = database.Owner.query().filter_by(owner_id=owner_id).first()
    if owner is None:
        owner = database.Owner(owner_id=owner_id)

    for resource, scope in [('ram',  'total'), ('ram',  'single'),
                            ('disk', 'total'), ('disk', 'single'),
                            ('vms',  'total'), ('vms',  'active')]:
        val = getattr(opts, resource+scope)
        if val is not None:
            setattr(owner, resource+'_quota_'+scope, val if val >= 0 else None)

    database.session.commit()
    print owner
    return 0

if __name__ == '__main__':
    sys.exit(main(sys.argv))