Limit user sign up

Author: daya0576

beaverhabits/app/auth.py

@@ -83,7 +83,7 @@ async def user_create(
                     )
                     return user
     except UserAlreadyExists:
-        raise
+        raise Exception("User already exists!")
 
 
 def user_logout() -> bool:

beaverhabits/app/crud.py

@@ -37,3 +37,12 @@ async def get_user_habit_list(user: User) -> HabitListModel | None:
         result = await session.execute(stmt)
         logger.info(f"[CRUD] User {user.id} habit list query")
         return result.scalar()
+
+
+async def get_user_count() -> int:
+    async with get_async_session_context() as session:
+        stmt = select(User)
+        result = await session.execute(stmt)
+        user_count = len(result.all())
+        logger.info(f"[CRUD] User count query: {user_count}")
+        return user_count

beaverhabits/configs.py

@@ -9,6 +9,8 @@
 logging.getLogger("niceGUI").setLevel(logging.INFO)
 dotenv.load_dotenv()
 
+USER_DATA_FOLDER = ".user"
+
 
 class StorageType(Enum):
     SESSION = "SESSION"
@@ -29,7 +31,8 @@ class Settings(BaseSettings):
 
     # Storage
     HABITS_STORAGE: StorageType = StorageType.USER_DATABASE
-    DATABASE_URL: str = "sqlite+aiosqlite:///./habits.db"
+    DATABASE_URL: str = f"sqlite+aiosqlite:///./{USER_DATA_FOLDER}/habits.db"
+    MAX_USER_COUNT: int = -1
 
     # Customization
     FIRST_DAY_OF_WEEK: int = calendar.MONDAY

beaverhabits/routes.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from fastapi import Depends, FastAPI, Request
+from fastapi import Depends, FastAPI, HTTPException, Request
 from fastapi.responses import RedirectResponse
 from fastapi.routing import APIRoute
 from nicegui import app, ui
@@ -13,6 +13,7 @@
     user_create_token,
 )
 from .app.db import User
+from .app.crud import get_user_count
 from .app.users import current_active_user
 from .configs import settings
 from .frontend.add_page import add_page_ui
@@ -113,37 +114,51 @@ async def try_login():
 
     with ui.card().classes("absolute-center shadow-none w-96"):
         email = ui.input("email").on("keydown.enter", try_login)
-        password = ui.input("password", password=True, password_toggle_button=True).on(
-            "keydown.enter", try_login
-        )
-        ui.button("Continue", on_click=try_login).props('padding="xs lg"')
-        ui.separator()
-        with ui.row():
-            ui.label("New around here?").classes("text-sm")
-            ui.link("Create account", target="/register").classes("text-sm")
+        email.classes("w-56")
+        
+        password = ui.input("password", password=True, password_toggle_button=True)
+        password.on("keydown.enter", try_login)
+        password.classes("w-56")
+        
+        with ui.element("div").classes("flex mt-4 justify-between items-center"):
+            ui.button("Continue", on_click=try_login).props('padding="xs lg"')
+
+        if not await get_user_count() >= settings.MAX_USER_COUNT > 0:
+            ui.separator()
+            with ui.row():
+                ui.label("New around here?").classes("text-sm")
+                ui.link("Create account", target="/register").classes("text-sm")
 
 
 @ui.page("/register")
 async def register():
+    async def validate_max_user_count():
+        if await get_user_count() >= settings.MAX_USER_COUNT > 0:
+            raise HTTPException(status_code=404, detail="User limit reached")
+
     async def try_register():
         try:
+            await validate_max_user_count()
             user = await user_create(email=email.value, password=password.value)
         except Exception as e:
-            ui.notify(str(e))
+            ui.notify(str(e), color="negative")
         else:
             token = user and await user_create_token(user)
             if token is not None:
                 app.storage.user.update({"auth_token": token})
                 ui.navigate.to(app.storage.user.get("referrer_path", "/"))
 
+    await validate_max_user_count()
+
     with ui.card().classes("absolute-center shadow-none w-96"):
-        email = ui.input("email").on("keydown.enter", try_register)
+        email = ui.input("email").on("keydown.enter", try_register).classes("w-56")
         password = ui.input("password", password=True, password_toggle_button=True).on(
             "keydown.enter", try_register
-        )
+        ).classes("w-56")
 
         with ui.element("div").classes("flex mt-4 justify-between items-center"):
-            ui.button("Register", on_click=try_register)
+            ui.button("Register", on_click=try_register).props('padding="xs lg"')
+
         ui.separator()
         with ui.row():
             ui.label("Already have an account?")

beaverhabits/storage/user_file.py

@@ -2,7 +2,9 @@
 from typing import Optional
 
 from nicegui.storage import PersistentDict
+
 from beaverhabits.app.db import User
+from beaverhabits.configs import USER_DATA_FOLDER
 from beaverhabits.storage.dict import DictHabitList
 from beaverhabits.storage.storage import UserStorage
 
@@ -18,8 +20,8 @@ async def get_user_habit_list(self, user: User) -> Optional[DictHabitList]:
 
     async def save_user_habit_list(self, user: User, habit_list: DictHabitList) -> None:
         d = self._get_persistent_dict(user)
-        d["data"] = habit_list.data
+        d[KEY_NAME] = habit_list.data
 
     def _get_persistent_dict(self, user: User) -> PersistentDict:
-        path = Path(f".user/{str(user.email)}.json")
+        path = Path(f"{USER_DATA_FOLDER}/{str(user.email)}.json")
         return PersistentDict(path, encoding="utf-8")