✨ Add user and delete user commands

components/website-backend/app/command.py

@@ -3,6 +3,8 @@
 import logging
 import sys
 
+from app.commands.add_user import add_n_users
+from app.commands.delete_user import delete_users
 from app.commands.dump_db import dump_db
 from app.commands.execute_sql import execute_sql_command
 from app.commands.init_db import init_db
@@ -114,6 +116,42 @@
     help="SQL command",
 )
 
+add_user_parser = subparsers.add_parser(
+    "add-users",
+    help="Add n users",
+)
+add_user_parser.add_argument(
+    "-n",
+    type=int,
+    help="Number of users",
+)
+add_user_parser.add_argument(
+    "-f",
+    "--file",
+    type=str,
+    help="File to save users",
+    default="users.txt",
+)
+add_user_parser.add_argument(
+    "--default-balance",
+    type=int,
+    help="Default balance",
+    default=10000,
+)
+
+delete_user_parser = subparsers.add_parser(
+    "delete-users",
+    help="Delete users",
+)
+delete_user_parser.add_argument(
+    "-f",
+    "--file",
+    type=str,
+    help="File with users",
+    default="users.txt",
+)
+
+
 PROMPT_MESSAGE = (
     "Are you sure you want to reset the database, this will delete all data? [y/N] "
 )
@@ -156,6 +194,10 @@ async def main(command: str) -> None:
             await load_db(args.input)
         case "execute":
             await execute_sql_command(args.command)
+        case "add-users":
+            await add_n_users(args.n, args.file, args.default_balance)
+        case "delete-users":
+            await delete_users(args.file)
 
 
 if __name__ == "__main__":  # pragma: no cover

components/website-backend/app/commands/add_user.py

@@ -0,0 +1,60 @@
+import logging
+import random
+import string
+
+from app.core.types import SecurityScopes
+from app.crud.crud_account import account as accounts
+from app.crud.redis.balance import crud_balance
+from app.dependencies import get_db, get_redis
+from app.schemas import account as account_schema
+
+logger = logging.getLogger("app.command")
+
+
+async def add_user(username: str, password: str, default_balance: int) -> None:
+    logger.info(f"Creating new user with username {username} and password {password}")
+
+    async with get_db.get_session() as session:
+        # Create account
+        account = await accounts.create(
+            db=session,
+            obj_in=account_schema.AccountCreate(
+                username=username,
+                password=password,
+                first_name="User",
+                last_name="User",
+            ),
+        )
+        logger.info(f"User {username} created")
+
+        updated_account = account_schema.AccountUpdate(
+            **account_schema.Account.model_validate(account).model_dump()
+        )
+        updated_account.enabled = True
+        updated_account.scope = SecurityScopes.USER
+        await accounts.update(
+            db=session,
+            db_obj=account,
+            obj_in=updated_account,
+        )
+        logger.info(f"User {username} activated")
+
+    async with get_redis.get_client() as client:
+        # Create balance
+        await crud_balance.set(client, username, default_balance)
+        logger.info(f"User {username} balance created with 10000")
+
+
+async def add_n_users(n: int, file: str, default_balance: int) -> None:
+    with open(file, "w", encoding="utf-8") as f:
+        for i in range(n):
+            username = f"user{i}"
+            random_password = "".join(
+                [random.choice(string.ascii_letters) for _ in range(16)]
+            )
+            await add_user(username, random_password, default_balance)
+            f.write(f"{username} {random_password}\n")
+
+            logger.info(f"User {username} created with password {random_password}")
+
+    logger.info(f"{n} users created")

components/website-backend/app/commands/delete_user.py

@@ -0,0 +1,36 @@
+import logging
+
+from app.crud.crud_account import account as accounts
+from app.crud.redis.balance import crud_balance
+from app.crud.redis.stock import crud_stock
+from app.dependencies import get_db, get_redis
+
+logger = logging.getLogger("app.command")
+
+
+async def delete_user(username: str) -> None:
+    logger.info(f"Deleting user {username}")
+
+    async with get_db.get_session() as session:
+        account_query = await accounts.query(db=session, username=username)
+        for account in account_query:
+            await accounts.delete(db=session, id=account.id)
+            logger.info(f"User {username} deleted")
+
+    async with get_redis.get_client() as client:
+        await crud_balance.delete(client, username)
+        await crud_stock.delete_all(client, username)
+        logger.info(f"User {username} balance deleted")
+
+
+async def delete_users(file: str) -> None:
+    with open(file, "r", encoding="utf-8") as f:
+        for line in f:
+            username, _ = line.split()
+
+            try:
+                await delete_user(username)
+            except Exception as e:
+                logger.error(f"Error deleting user {username}: {e}")
+
+    logger.info("Users deleted")