Merge pull request #161 from Delemangi/quotas

Roles & Quotas

backend/alembic/versions/548cdd92ed41_setup_roles.py

@@ -0,0 +1,30 @@
+"""Setup Roles
+
+Revision ID: 548cdd92ed41
+Revises: 0f18e18f9ae9
+Create Date: 2024-06-16 14:36:40.496965
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = '548cdd92ed41'
+down_revision: Union[str, None] = '0f18e18f9ae9'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column('user', 'quota')
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('user', sa.Column('quota', sa.INTEGER(), autoincrement=False, nullable=True))
+    # ### end Alembic commands ###

backend/app/auth/models.py

@@ -16,7 +16,6 @@ class User(Base):
     username = Column(String, nullable=False, unique=True)
     password = Column(String, nullable=False)
     avatar = Column(String, nullable=True)
-    quota = Column(Integer, nullable=True)
     timestamp = Column(DateTime(timezone=True), nullable=True)
 
     role_id = Column(UUID(as_uuid=True), ForeignKey("role.id"), nullable=True)
@@ -29,8 +28,14 @@ class User(Base):
     shared_files = relationship("Share", back_populates="user", lazy="selectin")
     code_2fa = Column(String, nullable=True)
 
-    def has_remaining_quota(self: Self) -> bool:
-        return bool(self.quota != 0)
+    def has_remaining_files_quota(self: Self) -> bool:
+        return self.role.quota_files is None or len(self.files) < self.role.quota_files
+
+    def has_remaining_size_quota(self: Self, size: int) -> bool:
+        return self.role.quota_size is None or self.get_used_space() + size < self.role.quota_size
+
+    def get_used_space(self: Self) -> int:
+        return sum(file.size for file in self.files)
 
 
 class Role(Base):

backend/app/auth/router.py

@@ -32,7 +32,7 @@ async def simple_test() -> str:
 
 @router.post("/create-test-user", response_model=str)
 async def test(session: Annotated[AsyncSession, Depends(get_async_session)]) -> str:
-    await create_user("a", "a", 30, session)
+    await create_user("a", "a", session)
 
     return "Created a test user"
 
@@ -81,7 +81,7 @@ async def logout(
 async def register(
     user_schema: User, session: Annotated[AsyncSession, Depends(get_async_session)]
 ) -> RequestStatus:
-    user = await create_user(user_schema.username, user_schema.password, 30, session)
+    user = await create_user(user_schema.username, user_schema.password, session)
     return RequestStatus(message=f"User {user.username} registered successfully")
 
 
@@ -121,6 +121,8 @@ async def fetch_user_data(
 ) -> UserMetadata:
     return UserMetadata(
         username=str(current_user.username),
-        quota=int(current_user.quota),
+        role=str(current_user.role.name),
+        files_quota=int(current_user.role.quota_files),
+        size_quota=int(current_user.role.quota_size),
         is_2fa_enabled=(current_user.code_2fa is not None),
     )

backend/app/auth/schemas.py

@@ -25,5 +25,7 @@ class Code2FA(BaseModel):
 
 class UserMetadata(BaseModel):
     username: str
-    quota: int
+    role: str
+    files_quota: int
+    size_quota: int
     is_2fa_enabled: bool

backend/app/auth/service.py

@@ -10,23 +10,49 @@
 from passlib.context import CryptContext
 from sqlalchemy import ColumnElement, delete, select, update
 from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import joinedload
 
 from ..database import get_async_session
 from .constants import ACCESS_TOKEN_EXPIRE_MINUTES, ALGORITHM
 from .exceptions import USERNAME_TAKEN_EXCEPTION
-from .models import LoggedInTokens, User
+from .models import LoggedInTokens, Role, User
 
 SECRET_KEY = os.getenv("JWT_SECRET_KEY", "SECRET")
 
 pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/auth/login")
 
 
-async def create_user(
-    username: str, plain_password: str, quota: int, session: AsyncSession
-) -> User:
+async def get_user_role(session: AsyncSession) -> Role:
+    role = await session.execute(select(Role).where(Role.name == "user"))
+
+    scalar_role = role.scalar_one_or_none()
+
+    if scalar_role is None:
+        raise ValueError("Role 'user' not found")
+
+    return scalar_role
+
+
+async def get_admin_role(session: AsyncSession) -> Role:
+    role = await session.execute(select(Role).where(Role.name == "admin"))
+
+    scalar_role = role.scalar_one_or_none()
+
+    if scalar_role is None:
+        raise ValueError("Role 'admin' not found")
+
+    return scalar_role
+
+
+async def create_user(username: str, plain_password: str, session: AsyncSession) -> User:
     password = pwd_context.hash(plain_password)
-    user = User(username=username, password=password, quota=quota)
+
+    user_role = (
+        await get_admin_role(session) if username == "admin" else await get_user_role(session)
+    )
+
+    user = User(username=username, password=password, role_id=user_role.id)
 
     existing_user = await get_user_by_username(username, session)
 
@@ -137,7 +163,9 @@ async def get_user_by_id(user_id: UUID, session: AsyncSession) -> User | None:
 
 async def get_user_by_username(username: str, session: AsyncSession) -> User | None:
     async with session:
-        users = await session.execute(select(User).filter(User.username == username))
+        users = await session.execute(
+            select(User).options(joinedload(User.role)).filter(User.username == username)
+        )
         return users.scalar_one_or_none()
 
 

backend/app/database.py

@@ -1,6 +1,8 @@
 import os
 from collections.abc import AsyncGenerator
+from datetime import datetime
 
+from sqlalchemy import select
 from sqlalchemy.ext.asyncio import (
     AsyncEngine,
     AsyncSession,
@@ -11,6 +13,8 @@
 from alembic import command
 from alembic.config import Config
 
+from .auth.models import Role
+
 SQLALCHEMY_DATABASE_URL = os.getenv(
     "DATABASE_URL",
     "postgresql+asyncpg://synthra:synthra@database:5432/synthra",
@@ -29,7 +33,7 @@ def get_engine(cls: type["DatabaseEngine"]) -> AsyncEngine:
 
 
 class AsyncSessionMaker:
-    _instance: None | async_sessionmaker[AsyncSession] = None
+    _instance: async_sessionmaker[AsyncSession] | None = None
 
     @classmethod
     def get_sessionmaker(cls: type["AsyncSessionMaker"]) -> async_sessionmaker[AsyncSession]:
@@ -51,3 +55,22 @@ async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
 def run_migrations() -> None:
     alembic_cfg = Config("alembic.ini")
     command.upgrade(alembic_cfg, "head")
+
+
+async def initialize_database() -> None:
+    async with AsyncSessionMaker.get_sessionmaker()() as session:
+        await initialize_roles(session)
+
+
+async def initialize_roles(session: AsyncSession) -> None:
+    existing_roles = await session.execute(select(Role).limit(1))
+    if existing_roles.scalars().first() is not None:
+        return
+
+    roles = [
+        Role(name="admin", quota_size=1000000000, quota_files=50, timestamp=datetime.now()),
+        Role(name="user", quota_size=100000000, quota_files=10, timestamp=datetime.now()),
+    ]
+
+    session.add_all(roles)
+    await session.commit()

backend/app/files/service.py

@@ -9,7 +9,7 @@
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 from fastapi import Depends, UploadFile
-from sqlalchemy import delete, select, update
+from sqlalchemy import delete, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from ..auth.dependencies import get_current_user
@@ -27,7 +27,10 @@ async def upload_file(
     is_shared: bool = False,
     password: str | None = None,
 ) -> str:
-    if not current_user.has_remaining_quota():
+    if not current_user.has_remaining_files_quota():
+        raise QUOTA_EXCEPTION
+
+    if not current_user.has_remaining_size_quota(file.size or 0):
         raise QUOTA_EXCEPTION
 
     file_path = f"{uuid.uuid4()}{file.filename}"
@@ -56,11 +59,6 @@ async def upload_file(
         )
 
         session.add(file_db)
-        update_statement = (
-            update(User).where(User.id == current_user.id).values(quota=User.quota - 1)
-        )
-
-        await session.execute(update_statement)
 
     return str(file_path)
 

backend/app/main.py

@@ -7,7 +7,7 @@
 from starlette.middleware.cors import CORSMiddleware
 
 from .auth.router import router as auth_router
-from .database import run_migrations
+from .database import initialize_database, run_migrations
 from .files.constants import FILE_PATH
 from .files.router import router as file_router
 from .jobs import schedule_jobs
@@ -19,30 +19,40 @@
 
 @asynccontextmanager
 async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
-    print("Starting...")
+    print("Starting...", flush=True)
 
-    print("Creating file storage directory...")
-    Path.mkdir(Path(FILE_PATH), exist_ok=True)
+    print("Creating file storage directory...", flush=True)
+    Path(FILE_PATH).mkdir(parents=True, exist_ok=True)
+    print("File storage directory created", flush=True)
 
-    print("Running migrations...")
+    print("Running migrations...", flush=True)
     run_migrations()
+    print("Migrations complete", flush=True)
 
-    print("Scheduling jobs...")
+    print("Initializing database...", flush=True)
+    await initialize_database()
+    print("Database initialized", flush=True)
+
+    print("Scheduling jobs...", flush=True)
     scheduler = schedule_jobs()
+    print("Jobs scheduled", flush=True)
 
-    print("Server started")
+    print("Server started", flush=True)
 
     yield
 
-    print("Shutting down...")
+    print("Shutting down...", flush=True)
+
+    print("Stopping jobs...", flush=True)
     scheduler.shutdown(wait=False)
+    print("Jobs stopped", flush=True)
 
-    print("Server stopped")
+    print("Server stopped", flush=True)
 
 
 def make_app() -> FastAPI:
     # set debug=True to enable verbose logging
-    app = FastAPI(lifespan=lifespan, debug=True)
+    app = FastAPI(lifespan=lifespan)
 
     # URL Normalizer Middleware
     app.add_middleware(SlashNormalizerMiddleware)

docker-compose.yaml

@@ -58,3 +58,5 @@ services:
     environment:
       PGADMIN_DEFAULT_EMAIL: [email protected]
       PGADMIN_DEFAULT_PASSWORD: test
+    volumes:
+      - ./pgadmin_data:/var/lib/pgadmin

frontend/src/lib/types/UserMetadata.ts

@@ -1,5 +1,7 @@
 export type UserMetadata = {
   username: string;
-  quota: number;
+  files_quota: number;
+  role: string;
+  size_quota: number;
   is_2fa_enabled: boolean;
 };

frontend/src/routes/user/account/+page.svelte

@@ -1,19 +1,20 @@
 <script lang="ts">
   import type { UserMetadata } from '$lib/types/UserMetadata';
-  import { Anchor, Button, Flex, Text } from '@svelteuidev/core';
+  import { Anchor, Button, Flex, Text, Title } from '@svelteuidev/core';
   import { isAxiosError } from 'axios';
   import { LockClosed, LockOpen2 } from 'radix-icons-svelte';
   import { onMount } from 'svelte';
   import { getUserMetadata } from '../../../server/auth';
+  import { getFilesForSpecifiedUser } from '../../../server/files';
   import { getPermanentToken } from '../../../server/sharex';
   import { generateShareXTemplate } from '../../../utils/functions';
 
-  let username: string | null = null;
   let user: UserMetadata | null = null;
+  let numberOfFiles: number | null = null;
+  let storageUsed: number | null = null;
 
   onMount(async () => {
     let accessToken = localStorage.getItem('accessToken');
-    username = localStorage.getItem('username');
 
     if (!accessToken) {
       window.location.href = '/auth/login';
@@ -22,6 +23,9 @@
 
     try {
       const response = await getUserMetadata(accessToken);
+      const files = await getFilesForSpecifiedUser(accessToken);
+      numberOfFiles = files.length;
+      storageUsed = files.reduce((acc, file) => acc + file.size, 0);
 
       user = response.data;
     } catch (error) {
@@ -60,10 +64,14 @@
   };
 </script>
 
-<Text size="lg" align="center">Hi, {username}!</Text>
-<br />
-
 <Flex justify="center" align="center" direction="column" gap="lg">
+  <Title>Account</Title>
+  <Text>Username: {user?.username}</Text>
+  <Text>Role: {user?.role.toUpperCase()}</Text>
+
+  <br />
+
+  <Title>2FA</Title>
   {#if user?.is_2fa_enabled}
     <Anchor href="/user/2fa?option=disable">
       <Flex justify="center" gap="md">
@@ -86,5 +94,16 @@
     </Anchor>
   {/if}
 
+  <br />
+
+  <Title>ShareX</Title>
   <Button on:click={generateShareXConfig}>Generate ShareX Configuration</Button>
+
+  <br />
+
+  <Title>Quotas</Title>
+  <Text>{numberOfFiles} / {user?.files_quota ?? 0} Files</Text>
+  <Text
+    >{storageUsed} / {user?.size_quota ?? 0} B ({((storageUsed ?? 0) / 1000 / 1000).toFixed(2)} MB)
+  </Text>
 </Flex>

frontend/src/routes/user/home/+page.svelte

@@ -47,6 +47,11 @@
         return;
       }
 
+      if (error.response?.status === 403) {
+        alert('Your session has expired or you have reached your quota.');
+        return;
+      }
+
       alert('An error occurred while uploading the file.');
     }
   };