CI-614 - Allow override of destination file names when uploading via SDK (#151)

* expose aws s3 client

* add optional file_path_map to dataset upload

* bump version

* test when file is omitted from map

* finish comment

* add example

* lint

* add helper function to flatten files

* lint

Author: nathanthorpe

cirro/clients/s3.py

@@ -38,6 +38,9 @@ def __init__(self, creds_getter: Callable[[], AWSCredentials], enable_additional
         self._upload_args = dict(ChecksumAlgorithm='SHA256') if enable_additional_checksum else dict()
         self._download_args = dict(ChecksumMode='ENABLED') if enable_additional_checksum else dict()
 
+    def get_aws_client(self):
+        return self._client
+
     def upload_file(self, file_path: Path, bucket: str, key: str):
         file_size = file_path.stat().st_size
         file_name = file_path.name

cirro/file_utils.py

@@ -2,7 +2,7 @@
 import random
 import time
 from pathlib import Path, PurePath
-from typing import List, Union
+from typing import List, Union, Dict
 
 from boto3.exceptions import S3UploadFailedError
 from botocore.exceptions import ConnectionError
@@ -35,6 +35,23 @@ def matches_glob(file: Union[File, str]):
     ]
 
 
+def generate_flattened_file_map(files: List[PathLike]) -> Dict[PathLike, str]:
+    """
+    Generates a mapping of file paths "flattened" to their base name.
+
+    Example:  data1/sample1.fastq.gz -> sample1.fastq.gz
+
+    Args:
+        files: List[PathLike]: List of file paths
+
+    Returns:
+        Dict[PathLike, str]: Mapping of file paths to their base name
+    """
+    return {
+        file: Path(file).name for file in files
+    }
+
+
 def _is_hidden_file(file_path: Path):
     # Remove hidden files from listing, desktop.ini .DS_Store, etc.
     if os.name == 'nt':
@@ -105,6 +122,7 @@ def get_files_stats(files: List[PathLike]) -> DirectoryStatistics:
 
 def upload_directory(directory: PathLike,
                      files: List[PathLike],
+                     file_path_map: Dict[PathLike, str],
                      s3_client: S3Client,
                      bucket: str,
                      prefix: str,
@@ -117,6 +135,7 @@ def upload_directory(directory: PathLike,
         directory (str|Path): Path to directory
         files (typing.List[str|Path]): List of paths to files within the directory
             must be the same type as directory.
+        file_path_map (typing.Dict[str|Path, str]): Map of file paths from source to destination
         s3_client (cirro.clients.S3Client): S3 client
         bucket (str): S3 bucket
         prefix (str): S3 prefix
@@ -132,7 +151,13 @@ def upload_directory(directory: PathLike,
         else:
             file_path = file
 
-        file_relative = file_path.relative_to(directory).as_posix()
+        # Check if is present in the file_path_map
+        # if it is, use the mapped value as the destination path
+        if file in file_path_map:
+            file_relative = file_path_map[file]
+        else:
+            file_relative = file_path.relative_to(directory).as_posix()
+
         key = f'{prefix}/{file_relative}'
         success = False
 

cirro/services/dataset.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Union
+from typing import List, Optional, Union, Dict
 
 from cirro_api_client.v1.api.datasets import get_datasets, get_dataset, import_public_dataset, upload_dataset, \
     update_dataset, delete_dataset, get_dataset_manifest
@@ -222,17 +222,53 @@ def upload_files(self,
                      project_id: str,
                      dataset_id: str,
                      directory: PathLike,
-                     files: List[PathLike]) -> None:
+                     files: List[PathLike] = None,
+                     file_path_map: Dict[PathLike, str] = None) -> None:
         """
         Uploads files to a given dataset from the specified directory.
 
+        All files must be relative to the specified directory.
+        If files need to be flattened, or you are sourcing files from multiple directories,
+        please include `file_path_map` or call this method multiple times.
+
         Args:
             project_id (str): ID of the Project
             dataset_id (str): ID of the Dataset
             directory (str|Path): Path to directory
             files (typing.List[str|Path]): List of paths to files within the directory,
                 must be the same type as directory.
+            file_path_map (typing.Dict[str|Path, str|Path]): Optional mapping of file paths to upload
+             from source path to destination path, used to "re-write" paths within the dataset.
+        ```python
+        from cirro.cirro_client import CirroApi
+        from cirro.file_utils import generate_flattened_file_map
+
+        cirro = CirroApi()
+
+        directory = "~/Downloads"
+        # Re-write file paths
+        file_map = {
+            "data1/file1.fastq.gz": "file1.fastq.gz",
+            "data2/file2.fastq.gz": "file2.fastq.gz",
+            "file3.fastq.gz": "new_file3.txt"
+        }
+
+        # Or you could automate the flattening
+        files = ["data1/file1.fastq.gz"]
+        file_map = generate_flattened_file_map(files)
+
+        cirro.datasets.upload_files(
+            project_id="project-id",
+            dataset_id="dataset-id",
+            directory=directory,
+            files=list(file_map.keys()),
+            file_path_map=file_map
+        )
+        ```
         """
+        if file_path_map is None:
+            file_path_map = {}
+
         dataset = self.get(project_id, dataset_id)
 
         access_context = FileAccessContext.upload_dataset(
@@ -244,7 +280,8 @@ def upload_files(self,
         self._file_service.upload_files(
             access_context=access_context,
             directory=directory,
-            files=files
+            files=files,
+            file_path_map=file_path_map
         )
 
     def download_files(

cirro/services/file.py

@@ -3,6 +3,7 @@
 from functools import partial
 from typing import List, Dict
 
+from botocore.client import BaseClient
 from cirro_api_client import CirroApiClient
 from cirro_api_client.v1.api.file import generate_project_file_access_token
 from cirro_api_client.v1.models import AWSCredentials, ProjectAccessType
@@ -69,6 +70,17 @@ def _get_project_read_credentials(self, access_context: FileAccessContext):
 
         return self._read_token_cache[project_id]
 
+    def get_aws_s3_client(self, access_context: FileAccessContext) -> BaseClient:
+        """
+        Gets the underlying AWS S3 client to perform operations on files
+
+        This is seeded with refreshable credentials from the access_context parameter
+
+        This may be used to perform advanced operations, such as CopyObject, S3 Select, etc.
+        """
+        s3_client = self._generate_s3_client(access_context)
+        return s3_client.get_aws_client()
+
     def get_file(self, file: File) -> bytes:
         """
         Gets the contents of a file
@@ -92,11 +104,7 @@ def get_file_from_path(self, access_context: FileAccessContext, file_path: str)
         Returns:
             The raw bytes of the file
         """
-
-        s3_client = S3Client(
-            partial(self.get_access_credentials, access_context),
-            self.enable_additional_checksum
-        )
+        s3_client = self._generate_s3_client(access_context)
 
         full_path = f'{access_context.prefix}/{file_path}'.lstrip('/')
 
@@ -113,11 +121,7 @@ def create_file(self, access_context: FileAccessContext, key: str,
             contents (str): Content of object
             content_type (str):
         """
-
-        s3_client = S3Client(
-            partial(self.get_access_credentials, access_context),
-            self.enable_additional_checksum
-        )
+        s3_client = self._generate_s3_client(access_context)
 
         s3_client.create_object(
             key=key,
@@ -129,7 +133,8 @@ def create_file(self, access_context: FileAccessContext, key: str,
     def upload_files(self,
                      access_context: FileAccessContext,
                      directory: PathLike,
-                     files: List[PathLike]) -> None:
+                     files: List[PathLike],
+                     file_path_map: Dict[PathLike, str]) -> None:
         """
         Uploads a list of files from the specified directory
 
@@ -138,19 +143,18 @@ def upload_files(self,
             directory (str|Path): Path to directory
             files (typing.List[str|Path]): List of paths to files within the directory
                 must be the same type as directory.
+            file_path_map (typing.Dict[str|Path, str]): Optional mapping of file paths to upload
+             from source path to destination path, used to "re-write" paths within the dataset.
         """
-
-        s3_client = S3Client(
-            partial(self.get_access_credentials, access_context),
-            self.enable_additional_checksum
-        )
+        s3_client = self._generate_s3_client(access_context)
 
         upload_directory(
-            directory,
-            files,
-            s3_client,
-            access_context.bucket,
-            access_context.prefix,
+            directory=directory,
+            files=files,
+            file_path_map=file_path_map,
+            s3_client=s3_client,
+            bucket=access_context.bucket,
+            prefix=access_context.prefix,
             max_retries=self.transfer_retries
         )
 
@@ -163,10 +167,7 @@ def download_files(self, access_context: FileAccessContext, directory: str, file
             directory (str): download location
             files (List[str]): relative path of files to download
         """
-        s3_client = S3Client(
-            partial(self.get_access_credentials, access_context),
-            self.enable_additional_checksum
-        )
+        s3_client = self._generate_s3_client(access_context)
 
         download_directory(
             directory,
@@ -176,6 +177,15 @@ def download_files(self, access_context: FileAccessContext, directory: str, file
             access_context.prefix
         )
 
+    def _generate_s3_client(self, access_context: FileAccessContext):
+        """
+        Generates the Cirro-S3 client to perform operations on files
+        """
+        return S3Client(
+            partial(self.get_access_credentials, access_context),
+            self.enable_additional_checksum
+        )
+
 
 class FileEnabledService(BaseService):
     """

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cirro"
-version = "1.5.0"
+version = "1.5.1"
 description = "CLI tool and SDK for interacting with the Cirro platform"
 authors = ["Cirro Bio <[email protected]>"]
 license = "MIT"

tests/test_file_utils.py

@@ -29,6 +29,7 @@ def test_upload_directory_pathlike(self):
         ]
         upload_directory(directory=test_path,
                          files=test_files,
+                         file_path_map={},
                          s3_client=self.mock_s3_client,
                          bucket=self.test_bucket,
                          prefix=self.test_prefix)
@@ -47,6 +48,7 @@ def test_upload_directory_string(self):
         ]
         upload_directory(directory=test_path,
                          files=test_files,
+                         file_path_map={},
                          s3_client=self.mock_s3_client,
                          bucket=self.test_bucket,
                          prefix=self.test_prefix)
@@ -70,6 +72,41 @@ def test_upload_directory_different_types(self):
         with self.assertRaises(ValueError):
             upload_directory(directory=test_path,
                              files=test_files,
+                             file_path_map={},
                              s3_client=self.mock_s3_client,
                              bucket=self.test_bucket,
-                             prefix=self.test_prefix)
+                             prefix=self.test_prefix)
+
+    def test_upload_directory_file_map_included(self):
+        test_path = 'data'
+        test_files = [
+            'file1.txt',
+            'folder1/file2.txt',
+            'folder1/unmapped.txt'
+        ]
+
+        file_path_map = {
+            'file1.txt': 'mapped_file1.txt',
+            'folder1/file2.txt': 'mapped_file2.txt'
+            # unmapped file3
+        }
+
+        upload_directory(directory=test_path,
+                         files=test_files,
+                         file_path_map=file_path_map,
+                         s3_client=self.mock_s3_client,
+                         bucket=self.test_bucket,
+                         prefix=self.test_prefix)
+
+        # Check that upload file was called with the mapped key
+        self.mock_s3_client.upload_file.assert_has_calls([
+            call(file_path=Path(test_path, test_files[0]),
+                 bucket=self.test_bucket,
+                 key=f'{self.test_prefix}/mapped_file1.txt'),
+            call(file_path=Path(test_path, test_files[1]),
+                 bucket=self.test_bucket,
+                 key=f'{self.test_prefix}/mapped_file2.txt'),
+            call(file_path=Path(test_path, test_files[2]),
+                 bucket=self.test_bucket,
+                 key=f'{self.test_prefix}/folder1/unmapped.txt')
+        ], any_order=True)