80 lines
3.2 KiB
Python
80 lines
3.2 KiB
Python
# fastapi_server/user_mgmt.py
|
|
import db_setup as db
|
|
import bcrypt
|
|
import asyncio
|
|
import aiomysql
|
|
|
|
import db_setup as db
|
|
|
|
async def fetch_user_info(arg: str, type: str):
|
|
"""Fetches user info.
|
|
This kind of query is supposed to either return ONE user or nothing.
|
|
"""
|
|
if db.mode == "mysql":
|
|
async with db.mysql_pool.acquire() as conn:
|
|
async with conn.cursor(aiomysql.DictCursor) as cur:
|
|
await cur.execute(f"SELECT * FROM users WHERE {type} = %s", (arg,))
|
|
#print(cur.description)
|
|
return await cur.fetchone()
|
|
elif db.mode == "sqlite":
|
|
async with db.file.connect() as conn:
|
|
async with conn.execute(f"SELECT * FROM users WHERE {type} = ?", (arg,)) as cursor:
|
|
return await cursor.fetchone()
|
|
return None
|
|
|
|
|
|
async def create_user(first_name: str, last_name: str, email: str, password: str):
|
|
"""Creates a new user."""
|
|
hashed_password = bcrypt.hashpw(password.encode(), bcrypt.gensalt())
|
|
if db.mode == "mysql":
|
|
async with db.mysql_pool.acquire() as conn:
|
|
async with conn.cursor() as cur:
|
|
await cur.execute("INSERT INTO users (`first_name`,`last_name`,`email`,`password`) VALUES (%s,%s,%s,%s)", (first_name, last_name, email, hashed_password.decode()))
|
|
await conn.commit()
|
|
return True
|
|
elif db.mode == "sqlite":
|
|
async with db.file.connect() as conn:
|
|
await conn.execute("INSERT INTO users (`first_name`,`last_name`,`email`,`password`) VALUES (?, ?, ?, ?)", (first_name, last_name, email, hashed_password.decode()))
|
|
await conn.commit()
|
|
return True
|
|
return False
|
|
|
|
|
|
async def login_user(email: str, password: str):
|
|
"""Logs in a user."""
|
|
user_data = await fetch_user_info(email, "email")
|
|
print(user_data)
|
|
if user_data is None:
|
|
return None # User doesn't exist
|
|
else:
|
|
if bcrypt.checkpw(password.encode(), user_data.get("password").encode()):
|
|
return user_data # Login successful
|
|
else:
|
|
return None # Incorrect password
|
|
|
|
|
|
async def edit_user(user_id: int, updates: dict):
|
|
"""Edits an existing user."""
|
|
if db.mode == "mysql":
|
|
placeholders = ", ".join(["%s"] * len(updates))
|
|
if updates.get("password") is not None: # there is a password update
|
|
updates["password"] = bcrypt.hashpw(updates["password"].encode(), bcrypt.gensalt()).decode()
|
|
values = list(updates.values())
|
|
query = f"UPDATE users SET {', '.join([f'{key} = %s' for key in updates.keys()])} WHERE id = %s"
|
|
values.append(user_id)
|
|
async with db.mysql_pool.acquire() as conn:
|
|
async with conn.cursor() as cur:
|
|
await cur.execute(query, values)
|
|
await conn.commit()
|
|
return True
|
|
elif db.mode == "sqlite":
|
|
placeholders = ", ".join(["?"] * len(updates))
|
|
values = list(updates.values())
|
|
values.append(user_id)
|
|
query = f"UPDATE users SET {', '.join([f'{key} = ?' for key in updates.keys()])} WHERE id = ?"
|
|
async with db.file.connect() as conn:
|
|
await conn.execute(query, values)
|
|
await conn.commit()
|
|
return True
|
|
return False
|