source: main/waeup.stress/trunk/src/waeup/stress/scripts.py @ 10371

"""Helpers when inside an instance context.
"""
import os
import sys
import transaction
from zope.app.appsetup import database, config
from zope.app.publication.zopepublication import ZopePublication
from zope.app.wsgi import config as wsgi_config
from waeup.kofa.app import University

#: The name under which any artificially created University instance
#: will be put into the root folder
APP_NAME = 'stress_app'

IMPORTERS = {
    'faculties.csv': 'facultyprocessor',
    'departments.csv': 'departmentprocessor',
    'courses.csv': 'courseprocessor',
    'certificates.csv': 'certificateprocessor',
    'certificatecourses.csv': 'certificatecourseprocessor',
    'applicantcontainers.csv': 'applicants container processor',
    'applicants.csv': 'applicantprocessor',
    }

def configure(instance_path=''):
    config_file = os.path.join(
        instance_path, 'parts', 'etc', 'site.zcml')
    print "WAEUP.STRESS: configuring ZODB from %s" % (
        config_file,)
    context = config(config_file)
    return context

def open_zodb(instance_path=''):
    """Open a ZODB and configure it using a site.zcml file.

    If `instance_path` is not given, we assume to be in the root of
    some Zope instance and compute relative paths to
    'parts/etc/site.zcml' and 'var/filestorage/Data.fs'.

    If `instance_path` is given, these paths are os.path.joined with
    it.

    Returns a tuple containing a database object, an open connection
    to this DB and the root folder.
    """
    zope_conf = os.path.join(
        instance_path, 'parts', 'etc', 'zope.conf')
    config_file = os.path.join(
        instance_path, 'parts', 'etc', 'site.zcml')
    print "WAEUP.STRESS: opening ZODB"
    zodb_file = os.path.join(
        instance_path, 'var', 'filestorage', 'Data.fs')

    print "WAEUP.STRESS: configuring ZODB from %s" % (
        zope_conf)

    db = wsgi_config(zope_conf)

    # The way to configure a ZODB if only a single ZODB has to be set up.
    #context = config(config_file)
    #db = database(zodb_file)

    conn = db.open()
    root = conn.root()
    root_folder = root.get(ZopePublication.root_name, None)
    return db, conn, root_folder

def close_zodb(db):
    print "WAEUP.STRESS: closing ZODB"
    transaction.commit() # make sure changes are saved
    db.close()
    return

def zodb(func):
    """A decorator that does a ZODB setup and calls the decorated
    function with the ZODB root folder as first argument.
    """
    def wrapped(*args, **kw):
        instance_path = kw.get('instance_path', '')
        db, conn, root_folder = open_zodb(instance_path)
        try:
            func(root_folder, *args, **kw)
        finally:
            close_zodb(db)
    return wrapped

@zodb
def create_app(root_folder):
    """Create a University instance named 'stress_app' in ZODB.

    The `root_folder` arg is injected by the decorator and represents
    the root folder of an opened, writable ZODB.
    """
    root_folder[APP_NAME] = University()
    print "WAEUP.STRESS.SCRIPTS: created '%s' in root folder." % APP_NAME
    return

@zodb
def import_files(root_folder, *args):
    from zope.component.hooks import setSite
    from zope.component import getUtility
    from waeup.kofa.interfaces import IBatchProcessor
    setSite(root_folder[APP_NAME])
    for path in args:
        if not os.path.isfile(path):
            print "WAEUP.STRESS.SCRIPTS: no such file: %s" % path
            continue
        basename = os.path.basename(path)
        util_name = IMPORTERS.get(basename, None)
        if util_name is None:
            print "WAEUP.STRESS.SCRIPTS: not a valid import filename: %s" % (
                basename)
            print "  valid names: %s" % IMPORTERS.keys()
            continue
        util = getUtility(IBatchProcessor, name=util_name)
        for line in open(path, 'rb'):
            break
        headerfields = line.strip().split(',')
        print "WAEUP.STRESS.SCRIPTS: importing %s" % os.path.basename(path)
        results = util.doImport(path, headerfields)
        print "WAEUP.STRESS.SCRIPTS: import done."
    return

def main(cmd, *args, **kw):
    if cmd == 'install_app':
        configure()
        create_app()
        return
    elif cmd == 'import_csv':
        configure()
        create_app()
        import_files(*args)
        return

if __name__ == '__main__':
    cmd = sys.argv[1]
    args = []
    if len(sys.argv) > 2:
        args = sys.argv[2:]
    main(cmd, *args)