74 lines
3.0 KiB
Python
74 lines
3.0 KiB
Python
# fastapi_server/user_mgmt.py
|
|
from . import db_setup as db
|
|
import bcrypt
|
|
import asyncio
|
|
from aiomysql import DictCursor
|
|
|
|
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.pool.acquire() as conn:
|
|
async with conn.cursor(DictCursor) as cur:
|
|
await cur.execute("SELECT * FROM users WHERE %s = %s", (type, arg))
|
|
return await cur.fetchone()
|
|
elif db.mode == "sqlite":
|
|
async with db.file.connect() as conn:
|
|
async with conn.execute("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.db.mode == "mysql":
|
|
async with db.db.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.db.mode == "sqlite":
|
|
async with db.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")
|
|
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))
|
|
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.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
|