futurehouse-client 0.4.1.dev95__py3-none-any.whl → 0.4.2__py3-none-any.whl

futurehouse_client/clients/data_storage_methods.py

@@ -15,6 +15,7 @@ import aiohttp
 import requests as requests_lib
 from google.resumable_media import requests as resumable_requests
 from httpx import AsyncClient, Client, HTTPStatusError, codes
+from lmi.utils import gather_with_concurrency
 from requests.adapters import HTTPAdapter
 from tenacity import (
     before_sleep_log,
@@ -26,12 +27,17 @@ from tqdm import tqdm
 from urllib3.util.retry import Retry
 
 from futurehouse_client.models.data_storage_methods import (
+    CreateDatasetPayload,
     DataStorageLocationPayload,
     DataStorageRequestPayload,
     DataStorageResponse,
     DirectoryManifest,
     ManifestEntry,
 )
+from futurehouse_client.models.rest import (
+    DataStorageSearchPayload,
+    SearchCriterion,
+)
 from futurehouse_client.utils.general import retry_if_connection_error
 
 # this is only required if they're using a yaml manifest
@@ -54,6 +60,7 @@ INITIATE_HEADERS = {
     "x-goog-resumable": "start",
     "Content-Length": "0",
 }
+DOWNLOAD_CONCURRENCY = 3
 
 
 def _should_ignore_file(
@@ -438,6 +445,10 @@ class DataStorageCreationError(DataStorageError):
     """Raised when there's an error creating a data storage entry."""
 
 
+class DataStorageRetrievalError(DataStorageError):
+    """Raised when there's an error retrieving a data storage entry."""
+
+
 class ProgressWrapper:
     """Common progress wrapper for file uploads."""
 
@@ -462,32 +473,30 @@ class ProgressWrapper:
         return self.file_obj.tell()
 
 
-class DataStorageMethods:
+class DataStorageMethods:  # pylint: disable=too-many-public-methods
     """Data storage methods for RestClient.
 
     This class contains methods for interacting with the data storage API endpoints.
     """
 
     # needed for mypy `NoReturn`
-    def _handle_http_errors(self, e: HTTPStatusError) -> NoReturn:
+    def _handle_http_errors(self, e: HTTPStatusError, operation: str) -> NoReturn:
         """Handle common HTTP errors for data storage operations."""
         if e.response.status_code == codes.FORBIDDEN:
-            raise DataStorageCreationError(
-                "Not authorized to create data storage entries"
+            raise DataStorageError(
+                f"Error {operation} data storage entry, not authorized"
             ) from e
         if e.response.status_code == codes.UNPROCESSABLE_ENTITY:
-            raise DataStorageCreationError(
-                f"Invalid request payload: {e.response.text}"
-            ) from e
-        raise DataStorageCreationError(
-            f"Error creating data storage entry: {e.response.status_code} - {e.response.text}"
+            raise DataStorageError(f"Invalid request payload: {e.response.text}") from e
+        raise DataStorageError(
+            f"Error {operation} data storage entry: {e.response.status_code} - {e.response.text}"
         ) from e
 
     def _validate_file_path(self, file_path: str | Path) -> Path:
         """Validate file path exists and return Path object."""
         file_path = Path(file_path)
         if not file_path.exists():
-            raise DataStorageCreationError(f"File or directory not found: {file_path}")
+            raise DataStorageError(f"File or directory not found: {file_path}")
         return file_path
 
     def _build_zip_path(self, name: str, path: str | None) -> str:
@@ -529,19 +538,24 @@ class DataStorageMethods:
             return extracted_items[0]
         return extract_dir
 
-    async def _adownload_from_gcs(self, signed_url: str) -> Path:
+    async def _adownload_from_gcs(
+        self, signed_url: str, file_name: str | None = None
+    ) -> Path:
         """Download file from GCS using signed URL and handle unzipping if needed.
 
         Args:
             signed_url: The signed URL to download from
+            file_name: The name of the file to download
 
         Returns:
             Path to the downloaded file (or unzipped directory if it was a zip)
         """
+        file_name = file_name or "downloaded_file"
+
         try:
             with tempfile.TemporaryDirectory() as temp_dir_str:
                 temp_dir = Path(temp_dir_str)
-                temp_file = temp_dir / "downloaded_file"
+                temp_file = temp_dir / file_name
 
                 async with self.async_client.stream("GET", signed_url) as response:
                     response.raise_for_status()
@@ -549,11 +563,11 @@ class DataStorageMethods:
                     content_disposition = response.headers.get(
                         "content-disposition", ""
                     )
-                    filename = "downloaded_file"
+                    filename = file_name
                     if "filename=" in content_disposition:
                         filename = content_disposition.split("filename=")[-1].strip('"')
 
-                    if filename != "downloaded_file":
+                    if filename != file_name:
                         temp_file = temp_dir / filename
 
                     async with aiofiles.open(temp_file, "wb") as f:
@@ -583,21 +597,23 @@ class DataStorageMethods:
                     return final_file
 
         except Exception as e:
-            raise DataStorageCreationError(f"Failed to download from GCS: {e}") from e
+            raise DataStorageError(f"Failed to download from GCS: {e}") from e
 
-    def _download_from_gcs(self, signed_url: str) -> Path:
+    def _download_from_gcs(self, signed_url: str, file_name: str | None = None) -> Path:
         """Download file from GCS using signed URL and handle unzipping if needed (sync version).
 
         Args:
             signed_url: The signed URL to download from
-
+            file_name: The name of the file to download
         Returns:
             Path to the downloaded file (or unzipped directory if it was a zip)
         """
+        file_name = file_name or "downloaded_file"
+
         try:
             with tempfile.TemporaryDirectory() as temp_dir_str:
                 temp_dir = Path(temp_dir_str)
-                temp_file = temp_dir / "downloaded_file"
+                temp_file = temp_dir / file_name
 
                 with requests_lib.get(signed_url, stream=True, timeout=30) as response:
                     response.raise_for_status()
@@ -605,11 +621,11 @@ class DataStorageMethods:
                     content_disposition = response.headers.get(
                         "content-disposition", ""
                     )
-                    filename = "downloaded_file"
+                    filename = file_name
                     if "filename=" in content_disposition:
                         filename = content_disposition.split("filename=")[-1].strip('"')
 
-                    if filename != "downloaded_file":
+                    if filename != file_name:
                         temp_file = temp_dir / filename
 
                     with open(temp_file, "wb") as f:
@@ -639,7 +655,7 @@ class DataStorageMethods:
                     return final_file
 
         except Exception as e:
-            raise DataStorageCreationError(f"Failed to download from GCS: {e}") from e
+            raise DataStorageError(f"Failed to download from GCS: {e}") from e
 
     # =====================================
 
@@ -676,7 +692,7 @@ class DataStorageMethods:
     ) -> DataStorageResponse:
         """Create data storage entry via API (sync version)."""
         response = self.client.post(
-            "/v0.1/data-storage",
+            "/v0.1/data-storage/data-entries",
             json=payload.model_dump(mode="json", exclude_none=True),
         )
         response.raise_for_status()
@@ -687,7 +703,7 @@ class DataStorageMethods:
     ) -> DataStorageResponse:
         """Create data storage entry via API (async version)."""
         response = await self.async_client.post(
-            "/v0.1/data-storage",
+            "/v0.1/data-storage/data-entries",
             json=payload.model_dump(mode="json", exclude_none=True),
         )
         response.raise_for_status()
@@ -761,6 +777,7 @@ class DataStorageMethods:
         path: str | None = None,
         ignore_patterns: list[str] | None = None,
         ignore_filename: str = ".gitignore",
+        project_id: UUID | None = None,
     ) -> DataStorageResponse:
         """Upload a directory as a single zip file collection.
 
@@ -771,6 +788,7 @@ class DataStorageMethods:
             path: Optional GCS path for the zip file
             ignore_patterns: List of patterns to ignore when zipping
             ignore_filename: Name of ignore file to read from directory
+            project_id: ID of the project this data storage entry belongs to
 
         Returns:
             DataStorageResponse for the uploaded zip file
@@ -790,6 +808,7 @@ class DataStorageMethods:
                 description=description,
                 path=zip_gcs_path,
                 is_collection=True,
+                project_id=project_id,
             )
 
             logger.debug(
@@ -797,24 +816,30 @@ class DataStorageMethods:
             )
             data_storage_response = self._create_data_storage_entry(payload)
 
-            if not data_storage_response.signed_url:
-                raise DataStorageCreationError("No signed URL returned for zip upload")
+            for storage_location in data_storage_response.storage_locations:
+                if not storage_location.storage_config.signed_url:
+                    raise DataStorageCreationError(
+                        "No signed URL returned for zip upload"
+                    )
 
-            with tqdm(
-                total=zip_size,
-                unit="B",
-                unit_scale=True,
-                unit_divisor=1024,
-                desc=f"Uploading {dir_path.name} (zipped)",
-                miniters=1,
-                mininterval=0.1,
-            ) as pbar:
-                _upload_file_with_progress(
-                    data_storage_response.signed_url, temp_zip_path, pbar, zip_size
-                )
+                with tqdm(
+                    total=zip_size,
+                    unit="B",
+                    unit_scale=True,
+                    unit_divisor=1024,
+                    desc=f"Uploading {dir_path.name} (zipped)",
+                    miniters=1,
+                    mininterval=0.1,
+                ) as pbar:
+                    _upload_file_with_progress(
+                        storage_location.storage_config.signed_url,
+                        temp_zip_path,
+                        pbar,
+                        zip_size,
+                    )
 
             status_response = self.client.patch(
-                f"/v0.1/data-storage/{data_storage_response.data_storage.id}",
+                f"/v0.1/data-storage/data-entries/{data_storage_response.data_storage.id}",
                 json={"status": "active"},
             )
             status_response.raise_for_status()
@@ -832,6 +857,7 @@ class DataStorageMethods:
         path: str | None = None,
         ignore_patterns: list[str] | None = None,
         ignore_filename: str = ".gitignore",
+        project_id: UUID | None = None,
     ) -> DataStorageResponse:
         """Asynchronously upload a directory as a single zip file.
 
@@ -842,6 +868,7 @@ class DataStorageMethods:
             path: Optional GCS path for the zip file
             ignore_patterns: List of patterns to ignore when zipping
             ignore_filename: Name of ignore file to read from directory
+            project_id: ID of the project this data storage entry belongs to
 
         Returns:
             DataStorageResponse for the uploaded zip file
@@ -861,28 +888,35 @@ class DataStorageMethods:
                 description=description,
                 path=zip_gcs_path,
                 is_collection=True,
+                project_id=project_id,
             )
 
             data_storage_response = await self._acreate_data_storage_entry(payload)
 
-            if not data_storage_response.signed_url:
-                raise DataStorageCreationError("No signed URL returned for zip upload")
+            for storage_location in data_storage_response.storage_locations:
+                if not storage_location.storage_config.signed_url:
+                    raise DataStorageCreationError(
+                        "No signed URL returned for zip upload"
+                    )
 
-            with tqdm(
-                total=zip_size,
-                unit="B",
-                unit_scale=True,
-                unit_divisor=1024,
-                desc=f"Uploading {dir_path.name} (zipped)",
-                miniters=1,
-                mininterval=0.1,
-            ) as pbar:
-                await _aupload_file_with_progress(
-                    data_storage_response.signed_url, temp_zip_path, pbar, zip_size
-                )
+                with tqdm(
+                    total=zip_size,
+                    unit="B",
+                    unit_scale=True,
+                    unit_divisor=1024,
+                    desc=f"Uploading {dir_path.name} (zipped)",
+                    miniters=1,
+                    mininterval=0.1,
+                ) as pbar:
+                    await _aupload_file_with_progress(
+                        storage_location.storage_config.signed_url,
+                        temp_zip_path,
+                        pbar,
+                        zip_size,
+                    )
 
             status_response = await self.async_client.patch(
-                f"/v0.1/data-storage/{data_storage_response.data_storage.id}",
+                f"/v0.1/data-storage/data-entries/{data_storage_response.data_storage.id}",
                 json={"status": "active"},
             )
             status_response.raise_for_status()
@@ -898,6 +932,7 @@ class DataStorageMethods:
         file_path: Path,
         description: str | None,
         path: str | None = None,
+        project_id: UUID | None = None,
     ) -> DataStorageResponse:
         """Upload a single file."""
         file_size = file_path.stat().st_size
@@ -918,6 +953,7 @@ class DataStorageMethods:
                     content=text_content,
                     path=path,
                     is_collection=False,
+                    project_id=project_id,
                 )
 
                 logger.debug("Sending file as text content")
@@ -934,6 +970,7 @@ class DataStorageMethods:
             description=description,
             path=path,
             is_collection=False,
+            project_id=project_id,
         )
 
         logger.debug(
@@ -942,30 +979,34 @@ class DataStorageMethods:
 
         data_storage_response = self._create_data_storage_entry(payload)
 
-        if not data_storage_response.signed_url:
-            raise DataStorageCreationError("No signed URL returned from server")
+        for storage_location in data_storage_response.storage_locations:
+            if not storage_location.storage_config.signed_url:
+                raise DataStorageCreationError("No signed URL returned from server")
 
-        with tqdm(
-            total=file_size,
-            unit="B",
-            unit_scale=True,
-            unit_divisor=1024,
-            desc=f"Uploading {file_path.name}",
-            miniters=1,
-            mininterval=0.1,
-        ) as pbar:
-            try:
-                _upload_file_with_progress(
-                    data_storage_response.signed_url, file_path, pbar, file_size
-                )
-                logger.debug("File upload to signed URL completed successfully")
-            except Exception as e:
-                logger.error(f"Failed to upload file to signed URL: {e}")
-                raise
+            with tqdm(
+                total=file_size,
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1024,
+                desc=f"Uploading {file_path.name}",
+                miniters=1,
+                mininterval=0.1,
+            ) as pbar:
+                try:
+                    _upload_file_with_progress(
+                        storage_location.storage_config.signed_url,
+                        file_path,
+                        pbar,
+                        file_size,
+                    )
+                    logger.debug("File upload to signed URL completed successfully")
+                except Exception as e:
+                    logger.error(f"Failed to upload file to signed URL: {e}")
+                    raise
 
         logger.debug("Updating data storage status to active")
         status_response = self.client.patch(
-            f"/v0.1/data-storage/{data_storage_response.data_storage.id}",
+            f"/v0.1/data-storage/data-entries/{data_storage_response.data_storage.id}",
             json={"status": "active"},
         )
         status_response.raise_for_status()
@@ -980,6 +1021,7 @@ class DataStorageMethods:
         description: str | None,
         path: str | None = None,
         dataset_id: UUID | None = None,
+        project_id: UUID | None = None,
     ) -> DataStorageResponse:
         """Asynchronously upload a single file."""
         file_size, text_payload = self._prepare_single_file_upload(
@@ -1000,28 +1042,33 @@ class DataStorageMethods:
             path=path,
             is_collection=False,
             dataset_id=dataset_id,
+            project_id=project_id,
         )
 
         data_storage_response = await self._acreate_data_storage_entry(payload)
 
-        if not data_storage_response.signed_url:
-            raise DataStorageCreationError("No signed URL returned from server")
+        for location in data_storage_response.storage_locations:
+            if not location.storage_config.signed_url:
+                raise DataStorageCreationError(
+                    f"No signed URL returned from server for location: {location.id}"
+                )
 
-        with tqdm(
-            total=file_size,
-            unit="B",
-            unit_scale=True,
-            unit_divisor=1024,
-            desc=f"Uploading {file_path.name}",
-            miniters=1,
-            mininterval=0.1,
-        ) as pbar:
-            await _aupload_file_with_progress(
-                data_storage_response.signed_url, file_path, pbar, file_size
-            )
+            with tqdm(
+                total=file_size,
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1024,
+                desc=f"Uploading {file_path.name}",
+                miniters=1,
+                mininterval=0.1,
+                leave=False,
+            ) as pbar:
+                await _aupload_file_with_progress(
+                    location.storage_config.signed_url, file_path, pbar, file_size
+                )
 
         status_response = await self.async_client.patch(
-            f"/v0.1/data-storage/{data_storage_response.data_storage.id}",
+            f"/v0.1/data-storage/data-entries/{data_storage_response.data_storage.id}",
             json={"status": "active"},
         )
         status_response.raise_for_status()
@@ -1036,6 +1083,7 @@ class DataStorageMethods:
         path: str | None,
         parent_id: UUID | None,
         dataset_id: UUID | None = None,
+        project_id: UUID | None = None,
     ) -> DataStorageResponse:
         """Upload a single file with a parent ID (sync version)."""
         file_size, text_payload = self._prepare_single_file_upload(
@@ -1046,6 +1094,7 @@ class DataStorageMethods:
             logger.debug("Sending file as text content with parent_id")
             text_payload.parent_id = parent_id
             text_payload.dataset_id = dataset_id
+            text_payload.project_id = project_id
             return self._create_data_storage_entry(text_payload)
 
         logger.debug(
@@ -1058,28 +1107,30 @@ class DataStorageMethods:
             is_collection=False,
             parent_id=parent_id,
             dataset_id=dataset_id,
+            project_id=project_id,
         )
         data_storage_response = self._create_data_storage_entry(payload)
 
-        if not data_storage_response.signed_url:
-            raise DataStorageCreationError("No signed URL returned from server")
+        for location in data_storage_response.storage_locations:
+            if not location.storage_config.signed_url:
+                raise DataStorageCreationError("No signed URL returned from server")
 
-        with tqdm(
-            total=file_size,
-            unit="B",
-            unit_scale=True,
-            unit_divisor=1024,
-            desc=f"Uploading {file_path.name}",
-            miniters=1,
-            mininterval=0.1,
-            leave=False,
-        ) as pbar:
-            _upload_file_with_progress(
-                data_storage_response.signed_url, file_path, pbar, file_size
-            )
+            with tqdm(
+                total=file_size,
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1024,
+                desc=f"Uploading {file_path.name}",
+                miniters=1,
+                mininterval=0.1,
+                leave=False,
+            ) as pbar:
+                _upload_file_with_progress(
+                    location.storage_config.signed_url, file_path, pbar, file_size
+                )
 
         status_response = self.client.patch(
-            f"/v0.1/data-storage/{data_storage_response.data_storage.id}",
+            f"/v0.1/data-storage/data-entries/{data_storage_response.data_storage.id}",
             json={"status": "active"},
         )
         status_response.raise_for_status()
@@ -1092,6 +1143,7 @@ class DataStorageMethods:
         dir_manifest: DirectoryManifest,
         current_parent_id: UUID,
         dataset_id: UUID | None = None,
+        project_id: UUID | None = None,
     ) -> DataStorageResponse | None:
         """Process a single file item for upload."""
         try:
@@ -1109,6 +1161,7 @@ class DataStorageMethods:
                 path=None,
                 parent_id=current_parent_id,
                 dataset_id=dataset_id,
+                project_id=project_id,
             )
         except Exception as e:
             logger.error(f"Failed to upload file {item}: {e}")
@@ -1126,6 +1179,7 @@ class DataStorageMethods:
         base_dir: Path | None = None,
         dir_manifest: DirectoryManifest | None = None,
         dataset_id: UUID | None = None,
+        project_id: UUID | None = None,
     ) -> list[DataStorageResponse]:
         """Upload a directory with single dataset and individual file storage entries."""
         responses = []
@@ -1141,6 +1195,7 @@ class DataStorageMethods:
                 parent_id=None,
                 dataset_id=None,
                 is_collection=False,
+                project_id=project_id,
             )
 
             dir_response = self._create_data_storage_entry(payload)
@@ -1182,6 +1237,7 @@ class DataStorageMethods:
                     parent_id=current_parent_id,
                     dataset_id=current_dataset_id,
                     is_collection=False,
+                    project_id=project_id,
                 )
                 subdir_response = self._create_data_storage_entry(subdir_payload)
                 responses.append(subdir_response)
@@ -1197,6 +1253,7 @@ class DataStorageMethods:
                     base_dir=base_dir,
                     dir_manifest=subdir_manifest,
                     dataset_id=current_dataset_id,
+                    project_id=project_id,
                 )
                 responses.extend(subdir_responses)
             elif item.is_file():
@@ -1247,6 +1304,7 @@ class DataStorageMethods:
         path: str | None,
         parent_id: UUID | None,
         dataset_id: UUID | None = None,
+        project_id: UUID | None = None,
     ) -> DataStorageResponse:
         """Asynchronously upload a single file with a parent ID."""
         file_size, text_payload = self._prepare_single_file_upload(
@@ -1257,6 +1315,7 @@ class DataStorageMethods:
             logger.debug("Sending file as text content with parent_id")
             text_payload.parent_id = parent_id
             text_payload.dataset_id = dataset_id
+            text_payload.project_id = project_id
             return await self._acreate_data_storage_entry(text_payload)
 
         logger.debug(
@@ -1269,10 +1328,13 @@ class DataStorageMethods:
             is_collection=False,
             parent_id=parent_id,
             dataset_id=dataset_id,
+            project_id=project_id,
         )
         data_storage_response = await self._acreate_data_storage_entry(payload)
 
-        if not data_storage_response.signed_url:
+        storage_location = data_storage_response.storage_locations[0]
+
+        if not storage_location.storage_config.signed_url:
             raise DataStorageCreationError("No signed URL returned from server")
 
         with tqdm(
@@ -1285,11 +1347,11 @@ class DataStorageMethods:
             mininterval=0.1,
         ) as pbar:
             await _aupload_file_with_progress(
-                data_storage_response.signed_url, file_path, pbar, file_size
+                storage_location.storage_config.signed_url, file_path, pbar, file_size
             )
 
         status_response = await self.async_client.patch(
-            f"/v0.1/data-storage/{data_storage_response.data_storage.id}",
+            f"/v0.1/data-storage/data-entries/{data_storage_response.data_storage.id}",
             json={"status": "active"},
         )
         status_response.raise_for_status()
@@ -1302,6 +1364,7 @@ class DataStorageMethods:
         dir_manifest: DirectoryManifest,
         current_parent_id: UUID,
         dataset_id: UUID | None = None,
+        project_id: UUID | None = None,
     ) -> DataStorageResponse | None:
         """Asynchronously process a single file item for upload."""
         try:
@@ -1319,6 +1382,7 @@ class DataStorageMethods:
                 path=None,
                 parent_id=current_parent_id,
                 dataset_id=dataset_id,
+                project_id=project_id,
             )
         except Exception as e:
             logger.error(f"Failed to upload file {item}: {e}")
@@ -1336,6 +1400,7 @@ class DataStorageMethods:
         base_dir: Path | None = None,
         dir_manifest: DirectoryManifest | None = None,
         dataset_id: UUID | None = None,
+        project_id: UUID | None = None,
     ) -> list[DataStorageResponse]:
         """Upload a directory with single dataset and individual file storage entries (async)."""
         responses = []
@@ -1352,6 +1417,7 @@ class DataStorageMethods:
                 parent_id=None,
                 dataset_id=None,
                 is_collection=False,
+                project_id=project_id,
             )
 
             dir_response = await self._acreate_data_storage_entry(payload)
@@ -1392,6 +1458,7 @@ class DataStorageMethods:
                     parent_id=current_parent_id,
                     dataset_id=current_dataset_id,
                     is_collection=False,
+                    project_id=project_id,
                 )
                 subdir_response = await self._acreate_data_storage_entry(subdir_payload)
                 responses.append(subdir_response)
@@ -1407,6 +1474,7 @@ class DataStorageMethods:
                     base_dir=base_dir,
                     dir_manifest=subdir_manifest,
                     dataset_id=current_dataset_id,
+                    project_id=project_id,
                 )
                 responses.extend(subdir_responses)
             elif item.is_file():
@@ -1443,6 +1511,7 @@ class DataStorageMethods:
         content: str,
         description: str | None = None,
         path: str | None = None,
+        project_id: UUID | None = None,
     ) -> DataStorageResponse:
         """Store content as a string in the data storage system.
 
@@ -1451,6 +1520,7 @@ class DataStorageMethods:
             content: Content to store as a string
             description: Optional description of the data storage entry
             path: Optional path for the data storage entry
+            project_id: ID of the project this data storage entry belongs to
 
         Returns:
             DataStorageResponse containing the created data storage entry and storage locations
@@ -1464,10 +1534,11 @@ class DataStorageMethods:
                 content=content,
                 description=description,
                 path=path,
+                project_id=project_id,
             )
             return self._create_data_storage_entry(payload)
         except HTTPStatusError as e:
-            self._handle_http_errors(e)
+            self._handle_http_errors(e, "creating")
         except Exception as e:
             raise DataStorageCreationError(
                 f"An unexpected error occurred: {e!r}"
@@ -1486,6 +1557,7 @@ class DataStorageMethods:
         description: str | None = None,
         path: str | None = None,
         dataset_id: UUID | None = None,
+        project_id: UUID | None = None,
     ) -> DataStorageResponse:
         """Asynchronously store content as a string in the data storage system.
 
@@ -1495,6 +1567,7 @@ class DataStorageMethods:
             description: Optional description of the data storage entry
             path: Optional path for the data storage entry
             dataset_id: Optional dataset ID to add entry to, or None to create new dataset
+            project_id: ID of the project this data storage entry belongs to
 
         Returns:
             DataStorageResponse containing the created data storage entry and storage locations
@@ -1509,10 +1582,11 @@ class DataStorageMethods:
                 description=description,
                 path=path,
                 dataset_id=dataset_id,
+                project_id=project_id,
             )
             return await self._acreate_data_storage_entry(payload)
         except HTTPStatusError as e:
-            self._handle_http_errors(e)
+            self._handle_http_errors(e, "creating")
         except Exception as e:
             raise DataStorageCreationError(
                 f"An unexpected error occurred: {e!r}"
@@ -1534,6 +1608,7 @@ class DataStorageMethods:
         manifest_filename: str | None = None,
         ignore_patterns: list[str] | None = None,
         ignore_filename: str = ".gitignore",
+        project_id: UUID | None = None,
     ) -> DataStorageResponse:
         """Store file or directory content in the data storage system.
 
@@ -1552,6 +1627,7 @@ class DataStorageMethods:
             manifest_filename: Name of manifest file
             ignore_patterns: List of patterns to ignore when zipping directories
             ignore_filename: Name of ignore file to read from directory (default: .gitignore)
+            project_id: ID of the project this data storage entry belongs to
 
         Returns:
             DataStorageResponse containing the final data storage entry
@@ -1564,7 +1640,13 @@ class DataStorageMethods:
         try:
             if file_path.is_dir() and as_collection:
                 return self._upload_data_directory(
-                    name, file_path, description, path, ignore_patterns, ignore_filename
+                    name,
+                    file_path,
+                    description,
+                    path,
+                    ignore_patterns,
+                    ignore_filename,
+                    project_id,
                 )
             if file_path.is_dir() and not as_collection:
                 responses = self._upload_directory_hierarchically(
@@ -1574,16 +1656,19 @@ class DataStorageMethods:
                     manifest_filename=manifest_filename,
                     ignore_patterns=ignore_patterns,
                     ignore_filename=ignore_filename,
+                    project_id=project_id,
                 )
                 if not responses:
                     raise DataStorageCreationError(
                         "No data storage entries were created"
                     )
                 return responses[0]
-            return self._upload_data_single_file(name, file_path, description, path)
+            return self._upload_data_single_file(
+                name, file_path, description, path, project_id
+            )
 
         except HTTPStatusError as e:
-            self._handle_http_errors(e)
+            self._handle_http_errors(e, "creating")
         except Exception as e:
             raise DataStorageCreationError(
                 f"An unexpected error occurred during file upload: {e!r}"
@@ -1606,6 +1691,7 @@ class DataStorageMethods:
         ignore_patterns: list[str] | None = None,
         ignore_filename: str = ".gitignore",
         dataset_id: UUID | None = None,
+        project_id: UUID | None = None,
     ) -> DataStorageResponse:
         """Asynchronously store file or directory content in the data storage system.
 
@@ -1620,6 +1706,7 @@ class DataStorageMethods:
             ignore_patterns: List of patterns to ignore when zipping.
             ignore_filename: Name of ignore file to read (default: .gitignore).
             dataset_id: Optional dataset ID to add entry to, or None to create new dataset.
+            project_id: ID of the project this data storage entry belongs to
 
         Returns:
             The `DataStorageResponse` for the created entry. For hierarchical uploads,
@@ -1637,6 +1724,7 @@ class DataStorageMethods:
                         path,
                         ignore_patterns,
                         ignore_filename,
+                        project_id,
                     )
                 responses = await self._aupload_directory_hierarchically(
                     name=name,
@@ -1646,6 +1734,7 @@ class DataStorageMethods:
                     ignore_patterns=ignore_patterns,
                     ignore_filename=ignore_filename,
                     dataset_id=dataset_id,
+                    project_id=project_id,
                 )
                 if not responses:
                     raise DataStorageCreationError(
@@ -1653,11 +1742,11 @@ class DataStorageMethods:
                     )
                 return responses[0]
             return await self._aupload_data_single_file(
-                name, file_path, description, path, dataset_id
+                name, file_path, description, path, dataset_id, project_id
             )
 
         except HTTPStatusError as e:
-            self._handle_http_errors(e)
+            self._handle_http_errors(e, "creating")
         except Exception as e:
             raise DataStorageCreationError(
                 f"An unexpected error occurred during async file upload: {e!r}"
@@ -1674,7 +1763,9 @@ class DataStorageMethods:
         name: str,
         existing_location: DataStorageLocationPayload,
         description: str | None = None,
+        as_collection: bool = False,
         path: str | None = None,
+        project_id: UUID | None = None,
     ) -> DataStorageResponse:
         """Store content as a string in the data storage system.
 
@@ -1682,7 +1773,11 @@ class DataStorageMethods:
             name: Name of the data storage entry
             existing_location: Describes the existing data source location to register
             description: Optional description of the data storage entry
+            as_collection: If uploading a directory, `True` creates a single storage entry for
+                the whole directory and multiple storage locations for each file, `False` assumes
+                you are uploading a single file.
             path: Optional path for the data storage entry
+            project_id: ID of the project this data storage entry belongs to
 
         Returns:
             DataStorageResponse containing the created data storage entry and storage locations
@@ -1696,14 +1791,17 @@ class DataStorageMethods:
                 description=description,
                 path=path,
                 existing_location=existing_location,
+                project_id=project_id,
+                is_collection=as_collection,
             )
             response = self.client.post(
-                "/v0.1/data-storage", json=payload.model_dump(exclude_none=True)
+                "/v0.1/data-storage/data-entries",
+                json=payload.model_dump(exclude_none=True),
             )
             response.raise_for_status()
             return DataStorageResponse.model_validate(response.json())
         except HTTPStatusError as e:
-            self._handle_http_errors(e)
+            self._handle_http_errors(e, "creating")
         except Exception as e:
             raise DataStorageCreationError(
                 f"An unexpected error occurred: {e!r}"
@@ -1719,8 +1817,10 @@ class DataStorageMethods:
         self,
         name: str,
         existing_location: DataStorageLocationPayload,
+        as_collection: bool = False,
         description: str | None = None,
         path: str | None = None,
+        project_id: UUID | None = None,
     ) -> DataStorageResponse:
         """Store content as a string in the data storage system.
 
@@ -1728,7 +1828,11 @@ class DataStorageMethods:
             name: Name of the data storage entry
             existing_location: Describes the existing data source location to register
             description: Optional description of the data storage entry
+            as_collection: If uploading a directory, `True` creates a single storage entry for
+                the whole directory and multiple storage locations for each file, `False` assumes
+                you are uploading a single file.
             path: Optional path for the data storage entry
+            project_id: ID of the project this data storage entry belongs to
 
         Returns:
             DataStorageResponse containing the created data storage entry and storage locations
@@ -1742,19 +1846,290 @@ class DataStorageMethods:
                 description=description,
                 path=path,
                 existing_location=existing_location,
+                project_id=project_id,
+                is_collection=as_collection,
             )
             response = await self.async_client.post(
-                "/v0.1/data-storage", json=payload.model_dump(exclude_none=True)
+                "/v0.1/data-storage/data-entries",
+                json=payload.model_dump(exclude_none=True),
             )
             response.raise_for_status()
             return DataStorageResponse.model_validate(response.json())
         except HTTPStatusError as e:
-            self._handle_http_errors(e)
+            self._handle_http_errors(e, "creating")
         except Exception as e:
             raise DataStorageCreationError(
                 f"An unexpected error occurred: {e!r}"
             ) from e
 
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    def search_data_storage(
+        self,
+        criteria: list[SearchCriterion] | None = None,
+        size: int = 10,
+    ) -> list[dict]:
+        """Search data storage objects using structured criteria.
+
+        Args:
+            criteria: List of search criteria (SearchCriterion objects with field, operator, value)
+            size: Number of results to return (1-100)
+
+        Returns:
+            List of search results with scores and data storage information
+
+        Raises:
+            DataStorageCreationError: If there's an error searching data storage entries
+
+        Example:
+            from futurehouse_client.models.rest import SearchCriterion, SearchOperator
+            criteria = [
+                SearchCriterion(field="name", operator=SearchOperator.CONTAINS, value="document"),
+                SearchCriterion(field="project_id", operator=SearchOperator.EQUALS, value="my-project-id"),
+                SearchCriterion(field="status", operator=SearchOperator.EQUALS, value="active"),
+            ]
+            results = client.search_data_storage(criteria=criteria, size=20)
+        """
+        try:
+            payload = DataStorageSearchPayload(
+                criteria=criteria or [],
+                size=max(1, min(100, size)),  # Clamp between 1-100
+            )
+
+            response = self.client.post(
+                "/v0.1/data-storage/search",
+                json=payload.model_dump(mode="json"),
+            )
+            response.raise_for_status()
+            return response.json()
+
+        except HTTPStatusError as e:
+            if e.response.status_code == codes.SERVICE_UNAVAILABLE:
+                raise DataStorageCreationError(
+                    "Search functionality is currently unavailable"
+                ) from e
+            self._handle_http_errors(e, "searching")
+        except Exception as e:
+            raise DataStorageCreationError(
+                f"An unexpected error occurred during search: {e!r}"
+            ) from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    async def asearch_data_storage(
+        self,
+        criteria: list[SearchCriterion] | None = None,
+        size: int = 10,
+    ) -> list[dict]:
+        """Asynchronously search data storage objects using structured criteria.
+
+        Args:
+            criteria: List of search criteria (SearchCriterion objects with field, operator, value)
+            size: Number of results to return (1-100)
+
+        Returns:
+            List of search results with scores and data storage information
+
+        Raises:
+            DataStorageCreationError: If there's an error searching data storage entries
+
+        Example:
+            from futurehouse_client.models.rest import SearchCriterion, SearchOperator
+            criteria = [
+                SearchCriterion(field="name", operator=SearchOperator.CONTAINS, value="document"),
+                SearchCriterion(field="project_id", operator=SearchOperator.EQUALS, value="my-project-id"),
+                SearchCriterion(field="status", operator=SearchOperator.EQUALS, value="active"),
+            ]
+            results = await client.asearch_data_storage(criteria=criteria, size=20)
+        """
+        try:
+            payload = DataStorageSearchPayload(
+                criteria=criteria or [],
+                size=max(1, min(100, size)),  # Clamp between 1-100
+            )
+
+            response = await self.async_client.post(
+                "/v0.1/data-storage/search",
+                json=payload.model_dump(mode="json"),
+            )
+            response.raise_for_status()
+            return response.json()
+
+        except HTTPStatusError as e:
+            if e.response.status_code == codes.SERVICE_UNAVAILABLE:
+                raise DataStorageCreationError(
+                    "Search functionality is currently unavailable"
+                ) from e
+            self._handle_http_errors(e, "searching")
+        except Exception as e:
+            raise DataStorageCreationError(
+                f"An unexpected error occurred during async search: {e!r}"
+            ) from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    def similarity_search_data_storage(
+        self,
+        embedding: list[float],
+        size: int = 10,
+        min_score: float = 0.7,
+        dataset_id: UUID | None = None,
+        tags: list[str] | None = None,
+        user_id: str | None = None,
+        project_id: str | None = None,
+    ) -> list[dict]:
+        """Search data storage objects using vector similarity.
+
+        Args:
+            embedding: Embedding vector for similarity search
+            size: Number of results to return (1-100)
+            min_score: Minimum similarity score (0.0-1.0)
+            dataset_id: Optional dataset ID filter
+            tags: Optional list of tags to filter by
+            user_id: Optional user ID filter (admin only)
+            project_id: Optional project ID filter
+
+        Returns:
+            List of search results with similarity scores and data storage information
+
+        Raises:
+            DataStorageCreationError: If there's an error performing similarity search
+        """
+        try:
+            # Validate inputs
+            if not embedding:
+                raise DataStorageCreationError("Embedding vector is required")
+
+            if not all(isinstance(x, int | float) for x in embedding):
+                raise DataStorageCreationError("Embedding must be a list of numbers")
+
+            size = max(1, min(100, size))  # Clamp between 1-100
+            min_score = max(0.0, min(1.0, min_score))  # Clamp between 0.0-1.0
+
+            # Build request payload
+            payload = {
+                "embedding": embedding,
+                "size": size,
+                "min_score": min_score,
+            }
+
+            # Add optional filters
+            if dataset_id is not None:
+                payload["dataset_id"] = str(dataset_id)
+            if tags is not None:
+                payload["tags"] = tags
+            if user_id is not None:
+                payload["user_id"] = user_id
+            if project_id is not None:
+                payload["project_id"] = project_id
+
+            response = self.client.post(
+                "/v0.1/data-storage/similarity-search", json=payload
+            )
+            response.raise_for_status()
+            return response.json()
+
+        except HTTPStatusError as e:
+            if e.response.status_code == codes.SERVICE_UNAVAILABLE:
+                raise DataStorageCreationError(
+                    "Similarity search functionality is currently unavailable"
+                ) from e
+            self._handle_http_errors(e, "performing similarity search")
+        except Exception as e:
+            raise DataStorageCreationError(
+                f"An unexpected error occurred during similarity search: {e!r}"
+            ) from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    async def asimilarity_search_data_storage(
+        self,
+        embedding: list[float],
+        size: int = 10,
+        min_score: float = 0.7,
+        dataset_id: UUID | None = None,
+        tags: list[str] | None = None,
+        user_id: str | None = None,
+        project_id: str | None = None,
+    ) -> list[dict]:
+        """Asynchronously search data storage objects using vector similarity.
+
+        Args:
+            embedding: Embedding vector for similarity search
+            size: Number of results to return (1-100)
+            min_score: Minimum similarity score (0.0-1.0)
+            dataset_id: Optional dataset ID filter
+            tags: Optional list of tags to filter by
+            user_id: Optional user ID filter (admin only)
+            project_id: Optional project ID filter
+
+        Returns:
+            List of search results with similarity scores and data storage information
+
+        Raises:
+            DataStorageCreationError: If there's an error performing similarity search
+        """
+        try:
+            # Validate inputs
+            if not embedding:
+                raise DataStorageCreationError("Embedding vector is required")
+
+            if not all(isinstance(x, int | float) for x in embedding):
+                raise DataStorageCreationError("Embedding must be a list of numbers")
+
+            size = max(1, min(100, size))  # Clamp between 1-100
+            min_score = max(0.0, min(1.0, min_score))  # Clamp between 0.0-1.0
+
+            # Build request payload
+            payload = {
+                "embedding": embedding,
+                "size": size,
+                "min_score": min_score,
+            }
+
+            # Add optional filters
+            if dataset_id is not None:
+                payload["dataset_id"] = str(dataset_id)
+            if tags is not None:
+                payload["tags"] = tags
+            if user_id is not None:
+                payload["user_id"] = user_id
+            if project_id is not None:
+                payload["project_id"] = project_id
+
+            response = await self.async_client.post(
+                "/v0.1/data-storage/similarity-search", json=payload
+            )
+            response.raise_for_status()
+            return response.json()
+
+        except HTTPStatusError as e:
+            if e.response.status_code == codes.SERVICE_UNAVAILABLE:
+                raise DataStorageCreationError(
+                    "Similarity search functionality is currently unavailable"
+                ) from e
+            self._handle_http_errors(e, "performing similarity search")
+        except Exception as e:
+            raise DataStorageCreationError(
+                f"An unexpected error occurred during async similarity search: {e!r}"
+            ) from e
+
     # TODO: EVERYTHING BELOW THIS LINE SHOULD BE MOVED TO FH_TOOLS REPO
     # =================================================
     @retry(
@@ -1766,7 +2141,7 @@ class DataStorageMethods:
     def fetch_data_from_storage(
         self,
         data_storage_id: UUID | None = None,
-    ) -> str | Path | None:
+    ) -> str | Path | list[Path] | None:
         """Fetch data from the storage system (sync version).
 
         Args:
@@ -1775,29 +2150,45 @@ class DataStorageMethods:
         Returns:
             For PG_TABLE storage: string content
             For GCS storage: Path to downloaded file (may be unzipped if it was a zip)
+            For multi-location entries: dict of location IDs to dicts with signed URL and file name
             None if not found or error occurred
         """
         if not data_storage_id:
-            raise DataStorageCreationError(
+            raise DataStorageRetrievalError(
                 "data_storage_id must be provided at this time"
             )
 
         try:
-            response = self.client.get(f"/v0.1/data-storage/{data_storage_id}")
+            response = self.client.get(
+                f"/v0.1/data-storage/data-entries/{data_storage_id}", timeout=100
+            )
             response.raise_for_status()
             result = DataStorageResponse.model_validate(response.json())
 
-            storage_type = result.storage_location.storage_config.storage_type
+            if len(result.storage_locations) > 1:
+                return [
+                    self._download_from_gcs(
+                        location.storage_config.signed_url or "",
+                        (location.storage_config.location or "").split("/")[-1],
+                    )
+                    for location in result.storage_locations
+                ]
+
+            # Most scenarios will only have one location
+            storage_location = result.storage_locations[0]
+            storage_type = storage_location.storage_config.storage_type
 
             if storage_type == "gcs":
-                if not result.signed_url:
-                    raise DataStorageCreationError(
+                if not storage_location.storage_config.signed_url:
+                    raise DataStorageRetrievalError(
                         "No signed URL available for GCS download"
                     )
 
-                return self._download_from_gcs(result.signed_url)
+                return self._download_from_gcs(
+                    storage_location.storage_config.signed_url
+                )
 
-            if storage_type == "raw_content":
+            if storage_type in {"raw_content", "pg_table"}:
                 content = result.data_storage.content
                 if content is None:
                     logger.warning(
@@ -1806,12 +2197,12 @@ class DataStorageMethods:
                     return None
                 return content
 
-            raise DataStorageCreationError(f"Unsupported storage type: {storage_type}")
+            raise DataStorageRetrievalError(f"Unsupported storage type: {storage_type}")
 
         except HTTPStatusError as e:
-            self._handle_http_errors(e)
+            self._handle_http_errors(e, "retrieving")
         except Exception as e:
-            raise DataStorageCreationError(
+            raise DataStorageRetrievalError(
                 f"An unexpected error occurred: {e!r}"
             ) from e
 
@@ -1824,7 +2215,7 @@ class DataStorageMethods:
     async def afetch_data_from_storage(
         self,
         data_storage_id: UUID | None = None,
-    ) -> str | Path | None:
+    ) -> str | Path | list[Path] | None:
         """Fetch data from the storage system.
 
         Args:
@@ -1833,31 +2224,48 @@ class DataStorageMethods:
         Returns:
             For PG_TABLE storage: string content
             For GCS storage: Path to downloaded file (may be unzipped if it was a zip)
+            For multi-location entries: dict of location IDs to dicts with signed URL and file name
             None if not found or error occurred
         """
         if not data_storage_id:
-            raise DataStorageCreationError(
+            raise DataStorageRetrievalError(
                 "data_storage_id must be provided at this time"
             )
 
         try:
             response = await self.async_client.get(
-                f"/v0.1/data-storage/{data_storage_id}"
+                f"/v0.1/data-storage/data-entries/{data_storage_id}", timeout=100
             )
             response.raise_for_status()
             result = DataStorageResponse.model_validate(response.json())
 
-            storage_type = result.storage_location.storage_config.storage_type
+            if len(result.storage_locations) > 1:
+                return await gather_with_concurrency(
+                    DOWNLOAD_CONCURRENCY,
+                    [
+                        self._adownload_from_gcs(
+                            location.storage_config.signed_url or "",
+                            (location.storage_config.location or "").split("/")[-1],
+                        )
+                        for location in result.storage_locations
+                    ],
+                )
+
+            # Most scenarios will only have one location
+            storage_location = result.storage_locations[0]
+            storage_type = storage_location.storage_config.storage_type
 
             if storage_type == "gcs":
-                if not result.signed_url:
-                    raise DataStorageCreationError(
+                if not storage_location.storage_config.signed_url:
+                    raise DataStorageRetrievalError(
                         "No signed URL available for GCS download"
                     )
 
-                return await self._adownload_from_gcs(result.signed_url)
+                return await self._adownload_from_gcs(
+                    storage_location.storage_config.signed_url
+                )
 
-            if storage_type == "raw_content":
+            if storage_type in {"raw_content", "pg_table"}:
                 content = result.data_storage.content
                 if content is None:
                     logger.warning(
@@ -1866,11 +2274,189 @@ class DataStorageMethods:
                     return None
                 return content
 
-            raise DataStorageCreationError(f"Unsupported storage type: {storage_type}")
+            raise DataStorageRetrievalError(f"Unsupported storage type: {storage_type}")
 
         except HTTPStatusError as e:
-            self._handle_http_errors(e)
+            self._handle_http_errors(e, "retrieving")
+        except Exception as e:
+            raise DataStorageRetrievalError(
+                f"An unexpected error occurred: {e!r}"
+            ) from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    async def acreate_dataset(
+        self,
+        name: str,
+        description: str | None = None,
+        dataset_id: UUID | None = None,
+    ):
+        try:
+            payload = CreateDatasetPayload(
+                name=name,
+                description=description,
+                id=dataset_id,
+            )
+            response = await self.async_client.post(
+                "/v0.1/data-storage/datasets",
+                json=payload.model_dump(exclude_none=True),
+            )
+            response.raise_for_status()
+            return CreateDatasetPayload.model_validate(response.json())
+        except HTTPStatusError as e:
+            self._handle_http_errors(e, "creating")
+        except Exception as e:
+            raise DataStorageCreationError(
+                f"An unexpected error occurred: {e!r}"
+            ) from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    def create_dataset(
+        self,
+        name: str,
+        description: str | None = None,
+        dataset_id: UUID | None = None,
+    ):
+        try:
+            payload = CreateDatasetPayload(
+                name=name,
+                description=description,
+                id=dataset_id,
+            )
+            response = self.client.post(
+                "/v0.1/data-storage/datasets",
+                json=payload.model_dump(exclude_none=True),
+            )
+            response.raise_for_status()
+            return CreateDatasetPayload.model_validate(response.json())
+        except HTTPStatusError as e:
+            self._handle_http_errors(e, "creating")
         except Exception as e:
             raise DataStorageCreationError(
                 f"An unexpected error occurred: {e!r}"
             ) from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    async def adelete_dataset(self, dataset_id: UUID):
+        """Delete a dataset.
+
+        Note: This will delete all data storage entries associated with the dataset.
+
+        Args:
+            dataset_id: ID of the dataset to delete
+
+        Raises:
+            DataStorageError: If there's an error deleting the dataset
+        """
+        try:
+            await self.async_client.delete(f"/v0.1/data-storage/datasets/{dataset_id}")
+        except HTTPStatusError as e:
+            self._handle_http_errors(e, "deleting")
+        except Exception as e:
+            raise DataStorageError(f"An unexpected error occurred: {e!r}") from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    def delete_dataset(self, dataset_id: UUID):
+        """Delete a dataset.
+
+        Note: This will delete all data storage entries associated with the dataset.
+
+        Args:
+            dataset_id: ID of the dataset to delete
+
+        Raises:
+            DataStorageError: If there's an error deleting the dataset
+        """
+        try:
+            self.client.delete(f"/v0.1/data-storage/datasets/{dataset_id}")
+        except HTTPStatusError as e:
+            self._handle_http_errors(e, "deleting")
+        except Exception as e:
+            raise DataStorageError(f"An unexpected error occurred: {e!r}") from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    async def aget_dataset(self, dataset_id: UUID):
+        try:
+            response = await self.async_client.get(
+                f"/v0.1/data-storage/datasets/{dataset_id}"
+            )
+            response.raise_for_status()
+
+            return response.json()
+        except HTTPStatusError as e:
+            self._handle_http_errors(e, "retrieving")
+        except Exception as e:
+            raise DataStorageError(f"An unexpected error occurred: {e!r}") from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    def get_dataset(self, dataset_id: UUID):
+        try:
+            response = self.client.get(f"/v0.1/data-storage/datasets/{dataset_id}")
+            response.raise_for_status()
+
+            return response.json()
+        except HTTPStatusError as e:
+            self._handle_http_errors(e, "retrieving")
+        except Exception as e:
+            raise DataStorageError(f"An unexpected error occurred: {e!r}") from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    async def adelete_data_storage_entry(self, data_storage_entry_id: UUID):
+        try:
+            await self.async_client.delete(
+                f"/v0.1/data-storage/data-entries/{data_storage_entry_id}"
+            )
+        except HTTPStatusError as e:
+            self._handle_http_errors(e, "deleting")
+        except Exception as e:
+            raise DataStorageError(f"An unexpected error occurred: {e!r}") from e
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, max=10),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    def delete_data_storage_entry(self, data_storage_entry_id: UUID):
+        try:
+            self.client.delete(
+                f"/v0.1/data-storage/data-entries/{data_storage_entry_id}"
+            )
+        except HTTPStatusError as e:
+            self._handle_http_errors(e, "deleting")
+        except Exception as e:
+            raise DataStorageError(f"An unexpected error occurred: {e!r}") from e