diff --git a/backend/btrixcloud/basecrawls.py b/backend/btrixcloud/basecrawls.py index 2def786597..26b6f5b265 100644 --- a/backend/btrixcloud/basecrawls.py +++ b/backend/btrixcloud/basecrawls.py @@ -2,7 +2,7 @@ import os from datetime import timedelta -from typing import Optional, List, Union, Dict, Any, Type, TYPE_CHECKING, cast +from typing import Optional, List, Union, Dict, Any, Type, TYPE_CHECKING, cast, Tuple from uuid import UUID import urllib.parse @@ -797,6 +797,34 @@ async def get_all_crawl_search_values( "firstSeeds": list(first_seeds), } + async def calculate_org_crawl_file_storage( + self, oid: UUID, type_: Optional[str] = None + ) -> Tuple[int, int, int]: + """Calculate and return total size of crawl files in org. + + Returns tuple of (total, crawls only, uploads only) + """ + total_size = 0 + crawls_size = 0 + uploads_size = 0 + + cursor = self.crawls.find({"oid": oid}) + async for crawl_dict in cursor: + files = crawl_dict.get("files", []) + type_ = crawl_dict.get("type") + + item_size = 0 + for file_ in files: + item_size += file_.get("size", 0) + + total_size += item_size + if type_ == "crawl": + crawls_size += item_size + if type_ == "upload": + uploads_size += item_size + + return total_size, crawls_size, uploads_size + # ============================================================================ def init_base_crawls_api(app, user_dep, *args): diff --git a/backend/btrixcloud/orgs.py b/backend/btrixcloud/orgs.py index 7f5660aa52..fd1b4f6b14 100644 --- a/backend/btrixcloud/orgs.py +++ b/backend/btrixcloud/orgs.py @@ -69,6 +69,7 @@ UpdatedResponse, AddedResponse, AddedResponseId, + SuccessResponse, OrgInviteResponse, OrgAcceptInviteResponse, OrgDeleteInviteResponse, @@ -1319,6 +1320,39 @@ async def delete_org_and_data( # Delete org await self.orgs.delete_one({"_id": org.id}) + async def recalculate_storage(self, org: Organization) -> dict[str, bool]: + """Recalculate org storage use""" + try: + total_crawl_size, crawl_size, upload_size = ( + await self.base_crawl_ops.calculate_org_crawl_file_storage( + org.id, + ) + ) + profile_size = await self.profile_ops.calculate_org_profile_file_storage( + org.id + ) + + org_size = total_crawl_size + profile_size + + await self.orgs.find_one_and_update( + {"_id": org.id}, + { + "$set": { + "bytesStored": org_size, + "bytesStoredCrawls": crawl_size, + "bytesStoredUploads": upload_size, + "bytesStoredProfiles": profile_size, + } + }, + ) + # pylint: disable=broad-exception-caught, raise-missing-from + except Exception as err: + raise HTTPException( + status_code=400, detail=f"Error calculating size: {err}" + ) + + return {"success": True} + # ============================================================================ # pylint: disable=too-many-statements, too-many-arguments @@ -1534,6 +1568,12 @@ async def set_role( return {"updated": True} + @router.post( + "/recalculate-storage", tags=["organizations"], response_model=SuccessResponse + ) + async def recalculate_org_storage(org: Organization = Depends(org_owner_dep)): + return await ops.recalculate_storage(org) + @router.post("/invite", tags=["invites"], response_model=OrgInviteResponse) async def invite_user_to_org( invite: InviteToOrgRequest, diff --git a/backend/btrixcloud/profiles.py b/backend/btrixcloud/profiles.py index 0d367ed698..0364f1b6d4 100644 --- a/backend/btrixcloud/profiles.py +++ b/backend/btrixcloud/profiles.py @@ -475,6 +475,18 @@ async def add_profile_file_replica( {"$push": {"resource.replicas": {"name": ref.name, "custom": ref.custom}}}, ) + async def calculate_org_profile_file_storage(self, oid: UUID) -> int: + """Calculate and return total size of profile files in org""" + total_size = 0 + + cursor = self.profiles.find({"oid": oid}) + async for profile_dict in cursor: + file_ = profile_dict.get("resource") + if file_: + total_size += file_.get("size", 0) + + return total_size + # ============================================================================ # pylint: disable=redefined-builtin,invalid-name,too-many-locals,too-many-arguments diff --git a/backend/test/test_z_delete_org.py b/backend/test/test_z_delete_org.py index 4a6d1da395..2de04be119 100644 --- a/backend/test/test_z_delete_org.py +++ b/backend/test/test_z_delete_org.py @@ -3,6 +3,29 @@ from .conftest import API_PREFIX +def test_recalculate_org_storage(admin_auth_headers, default_org_id): + # Prior to deleting org, ensure recalculating storage works now that + # resources of all types have been created. + r = requests.post( + f"{API_PREFIX}/orgs/{default_org_id}/recalculate-storage", + headers=admin_auth_headers, + ) + assert r.status_code == 200 + assert r.json()["success"] + + r = requests.get( + f"{API_PREFIX}/orgs/{default_org_id}", + headers=admin_auth_headers, + ) + assert r.status_code == 200 + data = r.json() + + assert data["bytesStored"] > 0 + assert data["bytesStoredCrawls"] > 0 + assert data["bytesStoredUploads"] > 0 + assert data["bytesStoredProfiles"] > 0 + + def test_delete_org_non_superadmin(crawler_auth_headers, default_org_id): # Assert that non-superadmin can't delete org r = requests.delete(