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

import csv
import argparse
import re
import string
from sys import exit

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

from models import Base, Student

# 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()

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

# add administrator
students = {'0': 'Professor'}

if args.csvfile:
    # add students from csv file if available
    try:
        csvreader = csv.DictReader(open(args.csvfile, encoding='iso-8859-1'), delimiter=';', quotechar='"', skipinitialspace=True)
    except EnvironmentError:
        print(f'Error: CSV file "{args.csvfile}" not found.')
        exit(1)
    students.update({s['N.º']: fix(s['Nome']) for s in csvreader})

elif args.demo:
    # add a few fake students
    students.update({
        '1915': 'Alan Turing',
        '1938': 'Donald Knuth',
        '1815': 'Ada Lovelace',
        '1969': 'Linus Torvalds',
        '1955': 'Tim Burners-Lee',
        '1916': 'Claude Shannon',
        '1903': 'John von Neumann',
    })

print(f'Generating {len(students)} bcrypt password hashes. This will take some time...')

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

    for num, name in students.items():
        print('.', end='', flush=True)
        pw = (args.pw or num).encode('utf-8')
        session.add(Student(id=num, name=name, password=bcrypt.hashpw(pw, bcrypt.gensalt())))
    print()

    n = session.query(Student).count()
    print(f'New database created: {args.db}\n{n} user(s) inserted:')

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

except Exception as e:
    print(f'Error: Database "{args.db}" already exists?')
    session.rollback()
    exit(1)

else:
    # --- end session ---
    session.commit()