init_db.py 1.77 KB
import sys
import os
import csv
from configparser import ConfigParser
from sqlalchemy import create_engine
from sqlalchemy.schema import CreateSchema
from sqlalchemy.orm import sessionmaker
from zope.sqlalchemy import register
import transaction
from ..models import (
    Base,
    Jalur,
    Status,
    )


my_registry = dict()


def get_file(filename):
    base_dir = os.path.split(__file__)[0]
    fullpath = os.path.join(base_dir, 'data', filename)
    return open(fullpath)


def append_csv(table, filename, keys):
    db_session = my_registry['db_session']
    with get_file(filename) as f:
        reader = csv.DictReader(f)
        filter_ = dict()
        for cf in reader:
            for key in keys:
                filter_[key] = cf[key]
            q = db_session.query(table).filter_by(**filter_)
            found = q.first()
            if found:
                continue
            row = table()
            for fieldname in cf:
                val = cf[fieldname]
                if not val:
                    continue
                setattr(row, fieldname, val)
            db_session.add(row)


def main(argv=sys.argv[1:]):
    conf_file = argv[0]
    conf = ConfigParser()
    conf.read(conf_file)
    db_url = conf.get('main', 'db_url')
    engine = create_engine(db_url, echo=True)
    schema_name = 'im'
    with engine.connect() as conn:
        if not conn.dialect.has_schema(conn, schema_name):
            sql = CreateSchema(schema_name)
            conn.execute(sql)
            conn.commit()
    Base.metadata.create_all(engine)
    factory = sessionmaker(bind=engine)
    my_registry['db_session'] = db_session = factory()
    register(db_session)
    with transaction.manager:
        append_csv(Jalur, 'jalur.csv', ['id'])
        append_csv(Status, 'status.csv', ['id'])