compute total size in datcore as well

Author: sanderegg

packages/models-library/src/models_library/api_schemas_datcore_adapter/datasets.py

@@ -1,13 +1,17 @@
 from datetime import datetime
 from enum import Enum, unique
 from pathlib import Path
+from typing import Annotated
 
-from pydantic import BaseModel, ByteSize
+from pydantic import BaseModel, ByteSize, Field
 
 
 class DatasetMetaData(BaseModel):
     id: str
     display_name: str
+    size: Annotated[
+        ByteSize | None, Field(description="Size of the dataset in bytes if available")
+    ]
 
 
 @unique

services/datcore-adapter/src/simcore_service_datcore_adapter/api/rest/datasets.py

@@ -56,6 +56,31 @@ async def list_datasets(
     return create_page(datasets, total=total, params=params)  # type: ignore[return-value]
 
 
+@router.get(
+    "/datasets/{dataset_id}",
+    status_code=status.HTTP_200_OK,
+    response_model=DatasetMetaData,
+)
+@cancel_on_disconnect
+async def get_dataset(
+    request: Request,
+    x_datcore_api_key: Annotated[str, Header(..., description="Datcore API Key")],
+    x_datcore_api_secret: Annotated[str, Header(..., description="Datcore API Secret")],
+    pennsieve_client: Annotated[PennsieveApiClient, Depends(get_pennsieve_api_client)],
+    params: Annotated[Params, Depends()],
+    dataset_id: str,
+) -> DatasetMetaData:
+    assert request  # nosec
+    raw_params: RawParams = resolve_params(params).to_raw_params()
+    assert raw_params.limit is not None  # nosec
+    assert raw_params.offset is not None  # nosec
+    return await pennsieve_client.get_dataset(
+        api_key=x_datcore_api_key,
+        api_secret=x_datcore_api_secret,
+        dataset_id=dataset_id,
+    )
+
+
 @router.get(
     "/datasets/{dataset_id}/files",
     summary="list top level files/folders in a dataset",

services/datcore-adapter/src/simcore_service_datcore_adapter/modules/pennsieve.py

@@ -14,6 +14,7 @@
     DataType,
     FileMetaData,
 )
+from pydantic import ByteSize
 from servicelib.logging_utils import log_context
 from servicelib.utils import logged_gather
 from starlette import status
@@ -81,9 +82,9 @@ class PennsieveAuthorizationHeaders(TypedDict):
     Authorization: str
 
 
-_TTL_CACHE_AUTHORIZATION_HEADERS_SECONDS: Final[
-    int
-] = 3530  # NOTE: observed while developing this code, pennsieve authorizes 3600 seconds, so we cache a bit less
+_TTL_CACHE_AUTHORIZATION_HEADERS_SECONDS: Final[int] = (
+    3530  # NOTE: observed while developing this code, pennsieve authorizes 3600 seconds, so we cache a bit less
+)
 
 ExpirationTimeSecs = int
 
@@ -346,12 +347,25 @@ async def list_datasets(
                 DatasetMetaData(
                     id=d["content"]["id"],
                     display_name=d["content"]["name"],
+                    size=ByteSize(d["storage"]) if d["storage"] > 0 else None,
                 )
                 for d in dataset_page["datasets"]
             ],
             dataset_page["totalCount"],
         )
 
+    async def get_dataset(
+        self, api_key: str, api_secret: str, dataset_id: str
+    ) -> DatasetMetaData:
+        dataset_pck = await self._get_dataset(api_key, api_secret, dataset_id)
+        return DatasetMetaData(
+            id=dataset_pck["content"]["id"],
+            display_name=dataset_pck["content"]["name"],
+            size=(
+                ByteSize(dataset_pck["storage"]) if dataset_pck["storage"] > 0 else None
+            ),
+        )
+
     async def list_packages_in_dataset(
         self,
         api_key: str,

services/storage/src/simcore_service_storage/datcore_dsm.py

@@ -5,6 +5,7 @@
 import arrow
 from fastapi import FastAPI
 from models_library.api_schemas_storage.storage_schemas import (
+    UNDEFINED_SIZE_TYPE,
     DatCoreCollectionName,
     DatCoreDatasetName,
     DatCorePackageName,
@@ -16,6 +17,7 @@
 from models_library.projects_nodes_io import LocationID, LocationName, StorageFileID
 from models_library.users import UserID
 from pydantic import AnyUrl, ByteSize, NonNegativeInt, TypeAdapter, ValidationError
+from servicelib.utils import limited_as_completed
 
 from .constants import DATCORE_ID, DATCORE_STR
 from .dsm_factory import BaseDataManager
@@ -191,12 +193,53 @@ async def compute_path_total_size(self, user_id: UserID, *, path: Path) -> ByteS
         """returns the total size of an arbitrary path"""
         api_token, api_secret = await self._get_datcore_tokens(user_id)
         api_token, api_secret = _check_api_credentials(api_token, api_secret)
+
+        # if this is a dataset we might have the size directly
+        with contextlib.suppress(ValidationError):
+            dataset_id = TypeAdapter(DatCoreDatasetName).validate_python(f"{path}")
+            _, dataset_size = await datcore_adapter.get_dataset(
+                self.app,
+                api_key=api_token,
+                api_secret=api_secret,
+                dataset_id=dataset_id,
+            )
+            if dataset_size is not None:
+                return dataset_size
+
+        # generic computation
         try:
-            paths = await self.list_paths(
-                user_id, file_filter=path, cursor=None, limit=1
+            paths, cursor, total_number = await self.list_paths(
+                user_id, file_filter=path, cursor=None, limit=50
             )
+            accumulated_size = ByteSize(0)
+
+            next_folders: list[PathMetaData] = []
+            for p in paths:
+                if p.file_meta_data is not None:
+                    # this is a file
+                    assert (
+                        p.file_meta_data.file_size is not UNDEFINED_SIZE_TYPE
+                    )  # nosec
+                    assert isinstance(p.file_meta_data.file_size, ByteSize)  # nosec
+                    accumulated_size = ByteSize(
+                        accumulated_size + p.file_meta_data.file_size
+                    )
+                else:
+                    next_folders.append(p)
+            async for sbfolder_size_future in limited_as_completed(
+                (
+                    self.compute_path_total_size(user_id, path=sub_folder.path)
+                    for sub_folder in next_folders
+                ),
+                limit=3,
+            ):
+                size = await sbfolder_size_future
+                accumulated_size = ByteSize(accumulated_size + size)
+
+            return accumulated_size
             if len(paths) == 0:
                 return ByteSize(0)
+
         except ValidationError:
             # invalid path
             return ByteSize(0)

services/storage/src/simcore_service_storage/modules/datcore_adapter/datcore_adapter.py

@@ -19,7 +19,7 @@
     DatCorePackageName,
 )
 from models_library.users import UserID
-from pydantic import AnyUrl, BaseModel, NonNegativeInt, TypeAdapter
+from pydantic import AnyUrl, BaseModel, ByteSize, NonNegativeInt, TypeAdapter
 from servicelib.fastapi.client_session import get_client_session
 from servicelib.utils import logged_gather
 
@@ -317,6 +317,31 @@ async def list_datasets(
     )
 
 
+async def get_dataset(
+    app: FastAPI,
+    *,
+    api_key: str,
+    api_secret: str,
+    dataset_id: DatCoreDatasetName,
+) -> tuple[DatasetMetaData, ByteSize | None]:
+    response = await request(
+        app,
+        api_key,
+        api_secret,
+        "GET",
+        f"/datasets/{dataset_id}",
+    )
+    assert isinstance(response, dict)  # nosec
+    datcore_dataset = DatCoreDatasetMetaData(**response)
+
+    return (
+        DatasetMetaData(
+            dataset_id=datcore_dataset.id, display_name=datcore_dataset.display_name
+        ),
+        datcore_dataset.size,
+    )
+
+
 async def get_file_download_presigned_link(
     app: FastAPI, api_key: str, api_secret: str, file_id: str
 ) -> AnyUrl: