diff --git a/src/db/schemas.py b/src/db/schemas.py index 072d4c3..424b9de 100644 --- a/src/db/schemas.py +++ b/src/db/schemas.py @@ -107,3 +107,11 @@ class ArchiveTrigger(BaseModel): public: bool = True group_id: Annotated[str, Len(min_length=1)] | None = None tags: set[Tag] | None = set() + +class Usage(BaseModel): + monthly_urls: int = 0 + monthly_mbs: int = 0 + total_sheets: int = 0 + +class UsageResponse(Usage): + groups: dict[str, Usage] \ No newline at end of file diff --git a/src/db/user_state.py b/src/db/user_state.py index 8b3ad54..5b81a15 100644 --- a/src/db/user_state.py +++ b/src/db/user_state.py @@ -6,6 +6,7 @@ from sqlalchemy import func from db import crud, models from datetime import datetime from shared.user_groups import GroupPermissions +from db.schemas import Usage, UsageResponse class UserState: @@ -127,7 +128,7 @@ class UserState: if not hasattr(self, '_max_archive_lifespan_months'): self._max_archive_lifespan_months = self._helper_for_grouping_max_numerical_permissions("max_archive_lifespan_months") return self._max_archive_lifespan_months - + @property def max_monthly_urls(self) -> int: if not hasattr(self, '_max_monthly_urls'): @@ -158,7 +159,7 @@ class UserState: if not hasattr(self, '_active'): self._active = bool(self.read or self.read_public or self.archive_url or self.archive_sheet) return self._active - + def _helper_for_grouping_max_numerical_permissions(self, permission_name: str) -> int: """ Iterates one of the numerical permissions where -1 means no restrictions and returns either -1 or the maximum value, defaults according to GroupPermissions @@ -211,36 +212,29 @@ class UserState: ).group_by(models.Archive.group_id).all() # merge the two queries - usage_by_group = { - (url.group_id or ""): { - "monthly_urls": url.url_count, - "monthly_mbs": int(url.total_bytes / 1024 / 1024), - "total_sheets": 0 - } + usage_by_group: Dict[str, Usage] = { + (url.group_id or ""): + Usage(monthly_urls=url.url_count, monthly_mbs=int(url.total_bytes / 1024 / 1024)) for url in urls_by_group } for group_id, sheet_count in sheets_by_group.items(): group_id = group_id or "" if group_id in usage_by_group: - usage_by_group[group_id]["total_sheets"] = sheet_count + usage_by_group[group_id].total_sheets = sheet_count else: - usage_by_group[group_id] = { - "monthly_urls": 0, - "monthly_mbs": 0, - "total_sheets": sheet_count - } + usage_by_group[group_id] = Usage(total_sheets=sheet_count) # calculate totals total_sheets = sum([sheet.sheet_count for sheet in user_sheets]) total_bytes = sum([url.total_bytes for url in urls_by_group]) total_urls = sum([url.url_count for url in urls_by_group]) - return { - "total_sheets": total_sheets, - "monthly_urls": total_urls, - "monthly_mbs": int(total_bytes / 1024 / 1024), - "groups": usage_by_group - } + return UsageResponse( + monthly_urls=total_urls, + monthly_mbs=int(total_bytes / 1024 / 1024), + total_sheets=total_sheets, + groups=usage_by_group + ) def has_quota_monthly_sheets(self, group_id: str) -> bool: """ @@ -256,7 +250,7 @@ class UserState: return True return user_sheets < sheet_quota - def has_quota_max_monthly_urls(self, group_id:str) -> bool: + def has_quota_max_monthly_urls(self, group_id: str) -> bool: """ checks if a user has reached their monthly url quota for a group, if global then group should be empty string """ @@ -277,7 +271,7 @@ class UserState: return user_urls < quota - def has_quota_max_monthly_mbs(self, group_id:str) -> bool: + def has_quota_max_monthly_mbs(self, group_id: str) -> bool: """ checks if a user has reached their monthly MBs quota for a group, if global then group should be empty string """ diff --git a/src/endpoints/default.py b/src/endpoints/default.py index d5f712f..557d7de 100644 --- a/src/endpoints/default.py +++ b/src/endpoints/default.py @@ -6,7 +6,8 @@ from sqlalchemy.orm import Session from core.config import VERSION, BREAKING_CHANGES from core.logging import log_error -from db import crud, schemas +from db import crud +from db.schemas import ActiveUser, UsageResponse from db.database import get_db_dependency from db.user_state import UserState from web.security import get_user_auth, bearer_security, get_user_state @@ -35,7 +36,7 @@ async def health(): @default_router.get("/user/active", summary="Check if the user is active and can use the tool.") async def active( user: UserState = Depends(get_user_state), -) -> schemas.ActiveUser: +) -> ActiveUser: return {"active": user.active} @@ -48,7 +49,7 @@ def get_user_permissions( @default_router.get("/user/usage", summary="Get the user's monthly URLs/MBs usage along with the total active sheets, breakdown by group.") def get_user_usage( user: UserState = Depends(get_user_state), -): +) -> UsageResponse: if not user.active: raise HTTPException(status_code=403, detail="User is not active.") return user.usage() @@ -56,5 +57,5 @@ def get_user_usage( @default_router.get('/favicon.ico', include_in_schema=False) -async def favicon(): +async def favicon() -> FileResponse: return FileResponse("static/favicon.ico")