diff --git a/initdb_from_csv.py b/initdb_from_csv.py index a819e82..d209fb0 100755 --- a/initdb_from_csv.py +++ b/initdb_from_csv.py @@ -3,12 +3,9 @@ import sqlite3 import csv -# from optparse import OptionParser import argparse -# from hashlib import sha512 import bcrypt import os -import sys import string import re @@ -16,52 +13,16 @@ import re def fixname(s): return string.capwords(re.sub('\(.*\)', '', s).strip()) +def genstudent(reader, pw=''): + for i, r in enumerate(reader): + print('\rInserting students into database... [{}]'.format(i+1), end='') + num = r['N.ยบ'] + name = fixname(r['Nome']) + pwhash = bcrypt.hashpw(pw.encode('utf-8'), bcrypt.gensalt()) + yield (num, name, pwhash) + print('\nDone.') -# --------- 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('--pw', default='', type=str, help='initial password') -argparser.add_argument('csvfile', type=str, help='CSV filename') -args = argparser.parse_args() - -# parser = OptionParser('usage: %prog [options] inputfile.csv') - -# parser.add_option('--db', dest='db_filename', default='students.db', -# help='database filename to create [default: %default]') - -# parser.add_option('--pw', dest='password', default='', -# help='initial password [default: %default]') - -# (options, args) = parser.parse_args() - -# if len(args) != 1: -# print('You must specify a CSV file to import.\nUse option -h for help.') -# sys.exit() - - -if os.path.exists(args.db): - print('Database already exists. Aborting...') - sys.exit(0) - -# -------- Parse CSV -------- -with open(args.csvfile, encoding='iso-8859-1') as csvfile: # SIIUE format - csvreader = csv.reader(csvfile, delimiter=';', quotechar='"') - next(csvreader) # ignore header - - rows = [] - for r in csvreader: - number = r[0] - name = fixname(r[1]) - password = bcrypt.hashpw(args.pw.encode('utf-8'), bcrypt.gensalt()) - rows.append((number, name, password)) - - # add professor (superuser) - number = '0' - name = 'Professor' - password = bcrypt.hashpw(args.pw.encode('utf-8'), bcrypt.gensalt()) - rows.append((number, name, password)) - -# ---- CREATE DATABASE AND INSERT DATA ---- +# ---- DATABASE SCHEMA ---- sql_cmd = '''PRAGMA foreign_keys = ON; CREATE TABLE students ( number TEXT PRIMARY KEY, @@ -85,6 +46,30 @@ sql_cmd = '''PRAGMA foreign_keys = ON; FOREIGN KEY(student_id) REFERENCES students(number) );''' -with sqlite3.connect(args.db) as c: - c.executescript(sql_cmd) - c.executemany('INSERT INTO students VALUES (?,?,?)', rows) +# --------- 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('--pw', default='', type=str, help='initial password') +argparser.add_argument('csvfile', type=str, help='CSV filename') +args = argparser.parse_args() + + +# -------- Parse CSV -------- +with open(args.csvfile, encoding='iso-8859-1') as csvfile: # SIIUE format + reader = csv.DictReader(csvfile, delimiter=';', quotechar='"') + db_exists = os.path.exists(args.db) + + with sqlite3.connect(args.db) as c: + if not db_exists: + print('Creating new database "{}"...'.format(args.db)) + c.executescript(sql_cmd) + pwhash = bcrypt.hashpw(pw.encode('utf-8'), bcrypt.gensalt()) + c.execute('INSERT INTO students VALUES (?,?,?)', ('0', 'Professor', pwhash)) + else: + print('Database "{}" already exists.'.format(args.db)) + + print('Warning: Passwords are generated using bcrypt wich takes a lot of time...') + try: + c.executemany('INSERT INTO students VALUES (?,?,?)', genstudent(reader, args.pw)) + except sqlite3.IntegrityError: + print('\rStudents already in the database. Aborting!!!') -- libgit2 0.21.2