"""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 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. """ print "WAEUP.STRESS: opening ZODB" zodb_file = os.path.join( instance_path, 'var', 'filestorage', 'Data.fs') 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)