initdb.py 3.03 KB
#!/usr/bin/env python3

import csv
import argparse
import re
import string
import sys

import bcrypt
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from models import Base, Student, Answer

# SIIUE names have alien strings like "(TE)" and are sometimes capitalized
# We remove them so that students dont keep asking what it means
def fix(name):
    return string.capwords(re.sub('\(.*\)', '', name).strip())

# ===========================================================================
#   Parse command line options
argparser = argparse.ArgumentParser(description='Create new database from a CSV file (SIIUE format)')
argparser.add_argument('--db', default='students.db', type=str, help='database filename')
argparser.add_argument('--demo', action='store_true', help='initialize database with a few fake students')
argparser.add_argument('--pw', default='', type=str, help='default password')
argparser.add_argument('csvfile', nargs='?', type=str, default='', help='CSV filename')
args = argparser.parse_args()

# ===========================================================================
hashed_pw = bcrypt.hashpw(args.pw.encode('utf-8'), bcrypt.gensalt())

engine = create_engine('sqlite:///{}'.format(args.db), echo=False)
Base.metadata.create_all(engine)  # Criate schema if needed
Session = sessionmaker(bind=engine)

# --- start session ---
try:
    session = Session()

    # add administrator
    session.add(Student(id='0', name='Professor', password=hashed_pw))

    # add students
    if args.csvfile:
        # from csv file if available
        try:
            csvreader = csv.DictReader(open(args.csvfile, encoding='iso-8859-1'), delimiter=';', quotechar='"', skipinitialspace=True)
        except EnvironmentError:
            print('Error: CSV file "{0}" not found.'.format(args.csvfile))
            session.rollback()
            sys.exit(1)
        else:
            session.add_all([Student(id=r['N.º'], name=fix(r['Nome']), password=hashed_pw) for r in csvreader])
    elif args.demo:
        # add a few fake students
        fakes = [
            ['1915', 'Alan Turing'],
            ['1938', 'Donald Knuth'],
            ['1815', 'Ada Lovelace'],
            ['1969', 'Linus Torvalds'],
            ['1955', 'Tim Burners-Lee'],
            ['1916', 'Claude Shannon'],
            ['1903', 'John von Neumann'],
        ]
        session.add_all([Student(id=i, name=name, password=hashed_pw) for i,name in fakes])

    session.commit()

except Exception:
    print('Error: Database already exists.')
    session.rollback()
    sys.exit(1)

else:
    n = session.query(Student).count()
    print('New database created: {0}\n{1} user(s) inserted:'.format(args.db, n))

    users = session.query(Student).order_by(Student.id).all()
    print('    {0:8} - {1} (administrator)'.format(users[0].id, users[0].name))
    if n > 1:
        print('    {0:8} - {1}'.format(users[1].id, users[1].name))
    if n > 3:
        print('    ...        ...')
    if n > 2:
        print('    {0:8} - {1}'.format(users[-1].id, users[-1].name))

# --- end session ---