enpi_api.l2.client.api.file_api

  1import os
  2import time
  3from pathlib import Path
  4from typing import Generator, Sequence
  5from urllib.parse import urlparse
  6from uuid import UUID
  7
  8from loguru import logger
  9from typing_extensions import assert_never
 10
 11from enpi_api.l1 import openapi_client
 12from enpi_api.l2.types.api_error import ApiError
 13from enpi_api.l2.types.execution import Execution
 14from enpi_api.l2.types.file import FederatedCredentials, File, FileId, FileStatus, OnCollisionAction
 15from enpi_api.l2.types.log import LogLevel
 16from enpi_api.l2.types.tag import Tag, TagId
 17from enpi_api.l2.types.task import TaskState
 18from enpi_api.l2.types.workflow import WorkflowExecutionTaskId
 19from enpi_api.l2.util.file import download_file, unique_temp_dir, upload_file_to_s3
 20from enpi_api.l2.util.tag import tags_to_api_payload
 21
 22
 23class NameEmpty(Exception):
 24    """Thrown when the name of a file is empty, which is not allowed."""
 25
 26    def __init__(self) -> None:
 27        """@private"""
 28        super().__init__("Name cannot be empty")
 29
 30
 31class S3UploadFailed(Exception):
 32    """Indicates that the upload to S3 failed."""
 33
 34    def __init__(self, file_path: str | Path, error: Exception):
 35        """@private"""
 36        super().__init__(f"Failed to upload file `{file_path}` to S3, error: {error}")
 37
 38
 39class FileApi:
 40    _inner_api_client: openapi_client.ApiClient
 41    _log_level: LogLevel
 42
 43    def __init__(self, inner_api_client: openapi_client.ApiClient, log_level: LogLevel):
 44        """@private"""
 45        self._inner_api_client = inner_api_client
 46        self._log_level = log_level
 47
 48    def get_files(self, filename: str | None = None) -> Generator[File, None, None]:
 49        """Get a generator through all available files in the platform.
 50
 51        Args:
 52            filename (str | None): Optional filename for search by case-insensitive substring matching
 53
 54        Returns:
 55            Generator[enpi_api.l2.types.file.File, None, None]: A generator through all files in the platform.
 56
 57        Raises:
 58            enpi_api.l2.types.api_error.ApiError: If API request fails.
 59
 60        Example:
 61
 62            ```python
 63            with EnpiApiClient() as enpi_client:
 64                for file in enpi_client.file_api.get_files():
 65                    print(file)
 66            ```
 67        """
 68
 69        logger.info("Getting a generator through all files")
 70
 71        file_api_instance = openapi_client.FileApi(self._inner_api_client)
 72
 73        # Fetch the first page, there is always a first page, it may be empty
 74        try:
 75            get_files_response = file_api_instance.get_files(filename=filename)
 76        except openapi_client.ApiException as e:
 77            raise ApiError(e)
 78
 79        # `files` and `cursor` get overwritten in the loop below when fetching a new page
 80        files = get_files_response.files
 81        cursor = get_files_response.cursor
 82
 83        while True:
 84            for file in files:
 85                yield File.from_raw(file)
 86
 87            # Check if we need to fetch a next page
 88            if cursor is None:
 89                logger.debug("No more pages of files")
 90                return  # No more pages
 91
 92            # We have a cursor, so we need to get a next page
 93            logger.debug("Fetching next page of files")
 94            try:
 95                get_files_response = file_api_instance.get_files(cursor=cursor, filename=filename)
 96            except openapi_client.ApiException as e:
 97                raise ApiError(e)
 98            files = get_files_response.files
 99            cursor = get_files_response.cursor
100
101    def get_file_by_id(self, file_id: FileId) -> File:
102        """Get a single file by its ID.
103
104        Args:
105            file_id (enpi_api.l2.types.file.FileId): The ID of the file to get.
106
107        Returns:
108            enpi_api.l2.types.file.File: The file, with all its metadata.
109
110        Raises:
111            enpi_api.l2.types.api_error.ApiError: If API request fails.
112
113        Example:
114
115            ```python
116            with EnpiApiClient() as enpi_client:
117                example_file_id = FileId("00000000-0000-0000-0000-000000000000")
118                file: File = enpi_client.file_api.get_file_by_id(file_id=example_file_id)
119            ```
120        """
121
122        logger.info(f"Getting file with ID `{file_id}`")
123
124        file_api_instance = openapi_client.FileApi(self._inner_api_client)
125
126        try:
127            get_file_response = file_api_instance.get_file(file_id=UUID(file_id))
128        except openapi_client.ApiException as e:
129            raise ApiError(e)
130
131        file = File.from_raw(get_file_response.file)
132
133        return file
134
135    def delete_file_by_id(self, file_id: FileId) -> None:
136        """Delete a single file by its ID.
137
138        This will remove the file from the ENPICOM Platform.
139
140        Args:
141            file_id (enpi_api.l2.types.file.FileId): The ID of the file to delete.
142
143        Raises:
144            enpi_api.l2.types.api_error.ApiError: If API request fails.
145
146        Example:
147
148            ```python
149            with EnpiApiClient() as enpi_client:
150                example_file_id = FileId("00000000-0000-0000-0000-000000000000")
151                enpi_client.file_api.delete_file_by_id(file_id=example_file_id))
152            ```
153        """
154
155        logger.info(f"Deleting file with ID `{file_id}`")
156
157        file_api_instance = openapi_client.FileApi(self._inner_api_client)
158
159        try:
160            file_api_instance.delete_file(file_id=UUID(file_id), body={})
161        except openapi_client.ApiException as e:
162            raise ApiError(e)
163
164        logger.info(f"File with ID `{file_id}` successfully deleted")
165
166    def upload_file(
167        self,
168        file_path: str | Path,
169        tags: Sequence[Tag] = (),
170        on_collision: OnCollisionAction = OnCollisionAction.ERROR,
171    ) -> Execution[File]:
172        """Upload a file to the platform.
173
174        Args:
175            file_path (str | Path): The path to the file to upload.
176            tags (Sequence[enpi_api.l2.types.tag.Tag]): The tags to add to the file.
177            on_collision (enpi_api.l2.types.file.OnCollisionAction): The action to take when uploading a file with the same name as an existing file.
178
179        Returns:
180            enpi_api.l2.types.execution.Execution[enpi_api.l2.types.file.File]: An awaitable that returns the uploaded file
181              or an existing one if the `OnCollisionAction` is set to `SKIP`.
182
183        Raises:
184            enpi_api.l2.client.api.file_api.NameEmpty: If the name of the file is empty.
185            enpi_api.l2.client.api.file_api.S3UploadFailed: If the upload to S3 failed.
186            enpi_api.l2.types.api_error.ApiError: If API request fails.
187
188        Example:
189
190            ```python
191            with EnpiApiClient() as enpi_client:
192                file: File = enpi_client.file_api.upload_file(file_path="path/to/file.txt").wait()
193            ```
194        """
195
196        logger.info(f"Uploading file with path `{file_path}`")
197
198        file_api_instance = openapi_client.FileApi(self._inner_api_client)
199
200        # Uploading a file is a two-step process:
201        # 1. Request an S3 temporary credentials used for upload
202        logger.debug("Requesting temporary upload credentials.")
203
204        name = os.path.basename(file_path)
205        name = name.strip()
206        if not name:
207            raise NameEmpty()
208
209        upload_file_request = openapi_client.UploadFileRequest(name=name, tags=tags_to_api_payload(tags), on_collision=on_collision)
210
211        try:
212            upload_file_response = file_api_instance.upload_file(upload_file_request)
213        except openapi_client.ApiException as e:
214            raise ApiError(e)
215
216        assert upload_file_response.id is not None, "upload_file_response.id must not be None"
217        file_id = FileId(upload_file_response.id)
218
219        # In the event that the file already exists, and we chose to SKIP, then we can return the existing file
220        if upload_file_response.credentials is None:
221            logger.info(f"File with name `{name}` already exists, and `on_collision` is set to `SKIP`")
222            return Execution(wait=lambda: self.get_file_by_id(file_id), check_execution_state=lambda: TaskState.SUCCEEDED)
223
224        s3_federated_credentials = FederatedCredentials.model_validate(
225            upload_file_response.credentials,
226            from_attributes=True,
227        )
228
229        # 2. Upload the file by using temporary credentials and boto3 client
230        logger.debug("Uploading the file.")
231        try:
232            upload_file_to_s3(file_path, s3_federated_credentials)
233            logger.debug("Post-processing the file.")
234        except Exception as err:
235            raise S3UploadFailed(file_path, err)
236
237        def wait() -> File:
238            # A file is not immediately usable after uploading, it needs to be processed first
239            # So before you can use a file you need to wait for it to be processed
240            self.wait_for_file_to_be_processed(file_id)
241
242            logger.success(f"File uploaded with ID `{file_id}`")
243
244            return self.get_file_by_id(file_id)
245
246        return Execution(wait=wait, check_execution_state=lambda: TaskState.SUCCEEDED)
247
248    def download_file_by_id(
249        self,
250        file_id: FileId,
251        output_directory: str | Path | None = None,
252        name: str | None = None,
253    ) -> Path:
254        """Download a single file by its ID into the specified directory.
255
256        Download a file from the platform to your local machine. The file will be saved in the specified directory with
257        the name of the file as it was in the ENPICOM Platform. Alternatively you can overwrite the name by providing one
258        yourself as the `name` argument.
259
260        Args:
261            file_id (enpi_api.l2.types.file.FileId): The ID of the file to download.
262            output_directory (str | Path): The directory to save the file to. If left empty, a temporary directory will be used.
263            name (str | None): The name of the file. If not provided, a name will be generated.
264
265        Returns:
266            Path: The path to the downloaded file.
267
268        Raises:
269            enpi_api.l2.client.api.file_api.NameEmpty: If the name of the file is empty.
270            enpi_api.l2.types.api_error.ApiError: If API request fails.
271
272        Example:
273
274            ```python
275            with EnpiApiClient() as enpi_client:
276                my_directory = f"/path/to/files"
277                example_file_id = FileId("00000000-0000-0000-0000-000000000000")
278
279                # Assume the file has the name `my_data.fastq`
280                full_file_path = enpi_client.file_api.download_file_by_id(
281                    file_id=example_file_id,
282                    directory=my_directory
283                )
284                # `full_file_path` will now be `/path/to/files/my_data.fastq`
285            ```
286        """
287
288        file_api_instance = openapi_client.FileApi(self._inner_api_client)
289
290        try:
291            download_file_response = file_api_instance.download_file(file_id=UUID(file_id))
292        except openapi_client.ApiException as e:
293            raise ApiError(e)
294
295        # If no name is provided, we parse the URL to get the file name
296        if name is None:
297            parsed_url = urlparse(download_file_response.download_url)
298            name = Path(parsed_url.path).name
299
300        if not name or name == "":
301            raise NameEmpty()
302
303        # Ensure that the directory exists
304        if output_directory is None:
305            output_directory = unique_temp_dir()
306
307        os.makedirs(output_directory, exist_ok=True)
308
309        full_path = os.path.join(output_directory, name)
310
311        logger.info(f"Downloading file with ID `{file_id}` to `{full_path}`")
312        downloaded_file_path = download_file(download_file_response.download_url, full_path)
313        logger.success(f"File with ID `{file_id}` successfully downloaded to `{downloaded_file_path}`")
314
315        return downloaded_file_path
316
317    def download_export_by_workflow_execution_task_id(
318        self, task_id: WorkflowExecutionTaskId, output_directory: str | Path | None = None, name: str | None = None
319    ) -> Path:
320        """Download a single file by its job ID to the specified directory.
321
322        Download an export from a job to your local machine. The export will be saved in the specified directory with
323        the name of the file as it was in the job. Alternatively you can overwrite the name by providing one
324        yourself as the `name` argument.
325
326        Args:
327            workflow_execution_id (enpi_api.l2.types.workflow.WorkflowExecutionId): The ID of the workflow execution to download an export from.
328            output_directory (str | Path | None): The directory to save the file to. If none is provided, a temporary
329              directory will be used.
330            name (str | None): The name of the file. If not provided, a name will be generated.
331
332        Returns:
333            Path: The path to the downloaded file.
334
335        Raises:
336            enpi_api.l2.client.api.file_api.NameEmpty: If the name of the file is empty.
337            enpi_api.l2.types.api_error.ApiError: If API request fails.
338
339        Example:
340
341            ```python
342            with EnpiApiClient() as enpi_client:
343                my_directory = f"/path/to/files"
344                example_task_id = TaskId(1234)
345
346                # Assume the file has the name `my_data.fastq`
347                full_file_path = enpi_client.file_api.download_export_by_workflow_execution_task_id(
348                    task_id=example_task_id,
349                    directory=my_directory
350                )
351                # `full_file_path` will now be `/path/to/files/my_data.fastq`
352            ```
353        """
354
355        file_api_instance = openapi_client.FileApi(self._inner_api_client)
356
357        try:
358            download_file_response = file_api_instance.download_export(job_id=task_id)
359        except openapi_client.ApiException as e:
360            raise ApiError(e)
361
362        # If no name is provided, we parse the URL to get the file name
363        if name is None:
364            parsed_url = urlparse(download_file_response.download_url)
365            name = Path(parsed_url.path).name
366
367        if not name or name == "":
368            raise NameEmpty()
369
370        # Ensure that the directory exists
371        if output_directory is None:
372            output_directory = unique_temp_dir()
373
374        os.makedirs(output_directory, exist_ok=True)
375
376        full_path = os.path.join(output_directory, name)
377
378        logger.info(f"Downloading export from task with ID `{task_id}` to `{full_path}`")
379        downloaded_file_path = download_file(download_file_response.download_url, full_path)
380        logger.success(f"Export from task with ID `{task_id}` successfully downloaded to `{downloaded_file_path}`")
381
382        return downloaded_file_path
383
384    def update_tags(self, file_id: FileId, tags: list[Tag]) -> None:
385        """Update the tags of a file.
386
387        Adds and updates the given tags to the file. If a tag is already present on the file, the value will be
388        overwritten with the given value for the same tag.
389
390        Args:
391            file_id (enpi_api.l2.types.file.FileId): The ID of the file to update.
392            tags (list[enpi_api.l2.types.tag.Tag]): The tags that will be updated or added if they are not already present.
393
394        Raises:
395            enpi_api.l2.types.api_error.ApiError: If API request fails.
396
397        Example:
398
399            ```python
400            with EnpiApiClient() as enpi_client:
401                enpi_client.file_api.update_tags(
402                    file_id=FileId("00000000-0000-0000-0000-000000000000"),
403                    tags=[
404                        Tag(id=TagId(FileTags.CampaignId), value="my new value"),
405                        Tag(id=TagId(FileTags.ProjectId), value="another value")
406                    ]
407                )
408            ```
409        """
410
411        logger.info(f"Updating tags for file with ID `{file_id}`")
412
413        file_api_instance = openapi_client.FileApi(self._inner_api_client)
414
415        update_file_tags_request = openapi_client.UpdateTagsRequest(tags=tags_to_api_payload(tags))
416
417        try:
418            file_api_instance.update_tags(file_id=UUID(file_id), update_tags_request=update_file_tags_request)
419        except openapi_client.ApiException as e:
420            raise ApiError(e)
421
422        logger.success(f"Tags updated for file with ID `{file_id}`")
423
424    def remove_tags(self, file_id: FileId, tags: list[TagId]) -> None:
425        """Remove the specified tags from a file.
426
427        Args:
428            file_id (enpi_api.l2.types.file.FileId): The ID of the file to update.
429            tags (List[enpi_api.l2.types.tag.TagId]): The tags that will be removed from the file.
430
431        Raises:
432            enpi_api.l2.types.api_error.ApiError: If API request fails.
433
434        Example:
435
436            ```python
437            with EnpiApiClient() as enpi_client:
438                enpi_client.file_api.remove_tags(
439                    file_id=FileId("00000000-0000-0000-0000-000000000000"),
440                    tags=[TagId(FileTags.CampaignId), TagId(FileTags.ProjectId)]
441                )
442            ```
443        """
444
445        logger.info(f"Removing tags: {tags} from file with ID `{file_id}`")
446
447        file_api_instance = openapi_client.FileApi(self._inner_api_client)
448
449        delete_tags_request = openapi_client.DeleteTagsRequest(tags=[int(x) for x in tags])
450
451        try:
452            file_api_instance.delete_tags(file_id=UUID(file_id), delete_tags_request=delete_tags_request)
453        except openapi_client.ApiException as e:
454            raise ApiError(e)
455
456        logger.success(f"Tags removed from file with ID `{file_id}`")
457
458    def wait_for_file_to_be_processed(self, file_id: FileId) -> None:
459        """Wait for a file to be processed.
460
461        Files are not immediately usable after uploading, they need to be processed first. This convenience method
462        waits for a file to be processed.
463
464        Args:
465            file_id (enpi_api.l2.types.file.FileId): The ID of the file to wait for.
466
467        Raises:
468            enpi_api.l2.types.api_error.ApiError: If API request fails.
469
470        Example:
471
472            ```python
473            with EnpiApiClient() as enpi_client:
474                enpi_client.file_api.wait_for_file_to_be_processed(file_id=FileId("00000000-0000-0000-0000-000000000000"))
475            ```
476        """
477
478        logger.info(f"Waiting for file with ID `{file_id}` to be processed")
479
480        poll_interval_seconds = 1
481
482        # We do not know how long it will take for a file to be processed, so we poll the file until it is processed
483        while True:
484            file = self.get_file_by_id(file_id)
485
486            if file.status == FileStatus.PROCESSED:
487                logger.success(f"File with ID `{file_id}` has been processed")
488                return
489            elif file.status == FileStatus.PROCESSING:
490                logger.debug(f"File with ID `{file_id}` is still being processed. Waiting for {poll_interval_seconds} seconds")
491                time.sleep(poll_interval_seconds)
492            else:
493                assert_never(file.status)