Commit 6d8bb39aa673c2a266d40d566be5ae113b465973

Authored by Miguel Barão
1 parent fecc60f6
Exists in master and in 1 other branch dev

Rewrite of the initdb_from_csv:

- now allows to add students to an existing database
- detects if students already exist
- show progress in the terminal during password generation
Showing 1 changed file with 36 additions and 51 deletions   Show diff stats
initdb_from_csv.py
@@ -3,12 +3,9 @@ @@ -3,12 +3,9 @@
3 3
4 import sqlite3 4 import sqlite3
5 import csv 5 import csv
6 -# from optparse import OptionParser  
7 import argparse 6 import argparse
8 -# from hashlib import sha512  
9 import bcrypt 7 import bcrypt
10 import os 8 import os
11 -import sys  
12 import string 9 import string
13 import re 10 import re
14 11
@@ -16,52 +13,16 @@ import re @@ -16,52 +13,16 @@ import re
16 def fixname(s): 13 def fixname(s):
17 return string.capwords(re.sub('\(.*\)', '', s).strip()) 14 return string.capwords(re.sub('\(.*\)', '', s).strip())
18 15
  16 +def genstudent(reader, pw=''):
  17 + for i, r in enumerate(reader):
  18 + print('\rInserting students into database... [{}]'.format(i+1), end='')
  19 + num = r['N.º']
  20 + name = fixname(r['Nome'])
  21 + pwhash = bcrypt.hashpw(pw.encode('utf-8'), bcrypt.gensalt())
  22 + yield (num, name, pwhash)
  23 + print('\nDone.')
19 24
20 -# --------- Parse command line options -----------  
21 -argparser = argparse.ArgumentParser(description='Create new database from a CSV file (SIIUE format)')  
22 -argparser.add_argument('--db', default='students.db', type=str, help='database filename')  
23 -argparser.add_argument('--pw', default='', type=str, help='initial password')  
24 -argparser.add_argument('csvfile', type=str, help='CSV filename')  
25 -args = argparser.parse_args()  
26 -  
27 -# parser = OptionParser('usage: %prog [options] inputfile.csv')  
28 -  
29 -# parser.add_option('--db', dest='db_filename', default='students.db',  
30 -# help='database filename to create [default: %default]')  
31 -  
32 -# parser.add_option('--pw', dest='password', default='',  
33 -# help='initial password [default: %default]')  
34 -  
35 -# (options, args) = parser.parse_args()  
36 -  
37 -# if len(args) != 1:  
38 -# print('You must specify a CSV file to import.\nUse option -h for help.')  
39 -# sys.exit()  
40 -  
41 -  
42 -if os.path.exists(args.db):  
43 - print('Database already exists. Aborting...')  
44 - sys.exit(0)  
45 -  
46 -# -------- Parse CSV --------  
47 -with open(args.csvfile, encoding='iso-8859-1') as csvfile: # SIIUE format  
48 - csvreader = csv.reader(csvfile, delimiter=';', quotechar='"')  
49 - next(csvreader) # ignore header  
50 -  
51 - rows = []  
52 - for r in csvreader:  
53 - number = r[0]  
54 - name = fixname(r[1])  
55 - password = bcrypt.hashpw(args.pw.encode('utf-8'), bcrypt.gensalt())  
56 - rows.append((number, name, password))  
57 -  
58 - # add professor (superuser)  
59 - number = '0'  
60 - name = 'Professor'  
61 - password = bcrypt.hashpw(args.pw.encode('utf-8'), bcrypt.gensalt())  
62 - rows.append((number, name, password))  
63 -  
64 -# ---- CREATE DATABASE AND INSERT DATA ---- 25 +# ---- DATABASE SCHEMA ----
65 sql_cmd = '''PRAGMA foreign_keys = ON; 26 sql_cmd = '''PRAGMA foreign_keys = ON;
66 CREATE TABLE students ( 27 CREATE TABLE students (
67 number TEXT PRIMARY KEY, 28 number TEXT PRIMARY KEY,
@@ -85,6 +46,30 @@ sql_cmd = '''PRAGMA foreign_keys = ON; @@ -85,6 +46,30 @@ sql_cmd = '''PRAGMA foreign_keys = ON;
85 FOREIGN KEY(student_id) REFERENCES students(number) 46 FOREIGN KEY(student_id) REFERENCES students(number)
86 );''' 47 );'''
87 48
88 -with sqlite3.connect(args.db) as c:  
89 - c.executescript(sql_cmd)  
90 - c.executemany('INSERT INTO students VALUES (?,?,?)', rows) 49 +# --------- Parse command line options -----------
  50 +argparser = argparse.ArgumentParser(description='Create new database from a CSV file (SIIUE format)')
  51 +argparser.add_argument('--db', default='students.db', type=str, help='database filename')
  52 +argparser.add_argument('--pw', default='', type=str, help='initial password')
  53 +argparser.add_argument('csvfile', type=str, help='CSV filename')
  54 +args = argparser.parse_args()
  55 +
  56 +
  57 +# -------- Parse CSV --------
  58 +with open(args.csvfile, encoding='iso-8859-1') as csvfile: # SIIUE format
  59 + reader = csv.DictReader(csvfile, delimiter=';', quotechar='"')
  60 + db_exists = os.path.exists(args.db)
  61 +
  62 + with sqlite3.connect(args.db) as c:
  63 + if not db_exists:
  64 + print('Creating new database "{}"...'.format(args.db))
  65 + c.executescript(sql_cmd)
  66 + pwhash = bcrypt.hashpw(pw.encode('utf-8'), bcrypt.gensalt())
  67 + c.execute('INSERT INTO students VALUES (?,?,?)', ('0', 'Professor', pwhash))
  68 + else:
  69 + print('Database "{}" already exists.'.format(args.db))
  70 +
  71 + print('Warning: Passwords are generated using bcrypt wich takes a lot of time...')
  72 + try:
  73 + c.executemany('INSERT INTO students VALUES (?,?,?)', genstudent(reader, args.pw))
  74 + except sqlite3.IntegrityError:
  75 + print('\rStudents already in the database. Aborting!!!')