Enforce uniqueness of users and move initialize_* to initialize.py

database.py

@@ -1,5 +1,6 @@
 import enum
 import random
+import threading
 import unicodedata
 import sqlite3
 
@@ -32,10 +33,12 @@ def connect():
 # Users
 # ------------------------------------------------------------------
 
+user_modify_lock = threading.Lock()
+
 def add_user(db, *, username, password, email, parent, status):
-	# TODO: Ensure users are unique
 	"""Add a user to the database
-	Will not commit the changes itself, so run .commit() on the database object yourself"""
+	Returns True is user was added succesfully and False if username was already in use
+	Will not commit the changes itself, so run .commit() on the database object yourself."""
 	global csprgn
 
 	assert type(username) == str
@@ -44,9 +47,6 @@ def add_user(db, *, username, password, email, parent, status):
 	assert type(parent) == int or parent is None
 	assert status in userstatus
 
-	# Generate a user ID. SQLite uses 64 bit signed ints, so generate at max 2⁶³-1
-	userid = csprng.randrange(2**63)
-
 	# Unicode normalize the username
 	username = unicodedata.normalize('NFKC', username)
 
@@ -57,10 +57,34 @@ def add_user(db, *, username, password, email, parent, status):
 	# Convert status into an int for storage
 	status = status.value
 
-	# Add the user into the database
-	cursor = db.cursor()
-	cursor.execute('PRAGMA foreign_keys = ON;') # Fail if we insert a user with bogus parent field
-	cursor.execute('INSERT INTO users VALUES (?, ?, ?, ?, ?, ?, ?);', (userid, parent, status, password, username, email, ''))
+	# We don't want any changes to the database to occur while we check if ID and username are available
+	with user_modify_lock:
+		cursor = db.cursor()
+
+		# Check that the username is unique
+		cursor.execute('SELECT id FROM users WHERE username = ?;', (username,))
+		results = cursor.fetchall()
+
+		if len(results) != 0:
+			return False # Username is already in use
+
+		# Generate a user ID
+		while True:
+			# SQLite uses 64 bit signed ints, so generate at max 2⁶³-1
+			userid = csprng.randrange(2**63)
+
+			# Check that the user ID is unique
+			cursor.execute('SELECT id FROM users WHERE id = ?;', (userid,))
+			results = cursor.fetchall()
+
+			if len(results) == 0:
+				break # It is unique
+
+		# Add the user into the database
+		cursor.execute('PRAGMA foreign_keys = ON;') # Fail if we insert a user with bogus parent field
+		cursor.execute('INSERT INTO users VALUES (?, ?, ?, ?, ?, ?, ?);', (userid, parent, status, password, username, email, ''))
+
+	return True
 
 def get_userid(db, username):
 	"""Returns the user ID associated with given username
@@ -121,32 +145,6 @@ def get_user_info(db, userid):
 
 	return UserInfo(userid, parent, status, username, email, comment)
 
-def initialize_users(db, admin_user, admin_password):
-	"""Creates a bare-bones user table with only admin user
-	This should never be run outside of the initialization script"""
-
-	cursor = db.cursor()
-
-	cursor.execute('''CREATE TABLE users (
-		id integer NOT NULL PRIMARY KEY,
-
-		parent integer,
-		status integer NOT NULL,
-
-		password text NOT NULL, 
-
-		username text NOT NULL,
-		email text NOT NULL,
-
-		comment text NOT NULL,
-
-		FOREIGN KEY(parent) REFERENCES users(id)
-	);''')
-
-	add_user(db, username = admin_user, password = admin_password, email = '', parent = None, status = userstatus.admin)
-
-	db.commit()
-
 # ------------------------------------------------------------------
 # Boards
 # ------------------------------------------------------------------
@@ -159,22 +157,3 @@ def list_boards(db):
 
 	# The results look like [('foo',), ('bar',), ('baz',)]
 	return [i[0] for i in results]
-
-def initialize_boards(db, boards):
-	"""Creates a table of boards
-	This should never be run outside of the initialization script"""
-
-	cursor = db.cursor()
-
-	cursor.execute('''CREATE TABLE boards (
-		id integer NOT NULL PRIMARY KEY,
-		name text NOT NULL
-	);''')
-
-	# .executemany() wants them in the format [("board1",), ("board2",), …]
-	boards = [(board_name,) for board_name in boards]
-
-	# Use NULL to have SQLite generate the IDs automatically
-	cursor.executemany('INSERT INTO boards VALUES (NULL, ?);', boards)
-
-	db.commit()

initialize.py

@@ -3,12 +3,58 @@ import sqlite3
 import config
 import database
 
+def initialize_users(db, admin_user, admin_password):
+	"""Creates a bare-bones user table with only admin user
+	This should never be run outside of the initialization script"""
+
+	cursor = db.cursor()
+
+	cursor.execute('''CREATE TABLE users (
+		id integer NOT NULL PRIMARY KEY,
+
+		parent integer,
+		status integer NOT NULL,
+
+		password text NOT NULL,
+
+		username text NOT NULL,
+		email text NOT NULL,
+
+		comment text NOT NULL,
+
+		FOREIGN KEY(parent) REFERENCES users(id)
+	);''')
+
+	# Since we just created the database, add_user cannot have conflicting user
+	database.add_user(db, username = admin_user, password = admin_password, email = '', parent = None, status = database.userstatus.admin)
+
+	db.commit()
+
+def initialize_boards(db, boards):
+	"""Creates a table of boards
+	This should never be run outside of the initialization script"""
+
+	cursor = db.cursor()
+
+	cursor.execute('''CREATE TABLE boards (
+		id integer NOT NULL PRIMARY KEY,
+		name text NOT NULL
+	);''')
+
+	# .executemany() wants them in the format [("board1",), ("board2",), …]
+	boards = [(board_name,) for board_name in boards]
+
+	# Use NULL to have SQLite generate the IDs automatically
+	cursor.executemany('INSERT INTO boards VALUES (NULL, ?);', boards)
+
+	db.commit()
+
 if __name__ == '__main__':
 	config.load('buranun.conf')
 	with database.connect() as db:
 		username = input('admin username: ')
 		password = input('admin password: ')
-		database.initialize_users(db, username, password)
+		initialize_users(db, username, password)
 
 		boards = input('boards: ').split()
-		database.initialize_boards(db, boards)
+		initialize_boards(db, boards)