Support prefix in helpers.get_all_files_in_bucket (#58)

* ?Support refix in helpers.get_all_files_in_bucket

* Fix param description

---------

Co-authored-by: John Kranz <[email protected]>

Author: berkeleysquare

ds3/ds3Helpers.py

@@ -420,8 +420,12 @@ def get_blob(self, bucket: str, get_object: HelperGetObject, offset: int, job_id
         stream.close()
         return get_object.object_name, offset
 
-    def get_all_files_in_bucket(self, destination_dir: str, bucket: str, objects_per_bp_job: int = 1000,
-                                max_threads: int = 5, job_name: str = None) -> List[str]:
+    def get_all_files_in_bucket(self, destination_dir: str,
+                                bucket: str,
+                                prefix: str = None,
+                                objects_per_bp_job: int = 1000,
+                                max_threads: int = 5,
+                                job_name: str = None) -> List[str]:
         """
         Retrieves all objects from a Black Pearl bucket.
 
@@ -439,6 +443,8 @@ def get_all_files_in_bucket(self, destination_dir: str, bucket: str, objects_per
             The number of concurrent objects being transferred at once (default 5).
         job_name : str
             The name to give the BP get jobs. All BP jobs that are created will have the same name.
+        prefix : str
+            Limits the response to objects that begin with the specified prefix.
         """
         truncated: str = 'true'
         marker = ""
@@ -447,6 +453,7 @@ def get_all_files_in_bucket(self, destination_dir: str, bucket: str, objects_per
             list_bucket = self.client.get_bucket(GetBucketRequest(bucket_name=bucket,
                                                                   max_keys=objects_per_bp_job,
                                                                   versions=False,
+                                                                  prefix=prefix,
                                                                   marker=marker))
 
             get_objects: List[HelperGetObject] = []

tests/helpersTests.py

@@ -163,7 +163,7 @@ def test_put_and_get_objects(self):
         destination.cleanup()
         client.delete_bucket_spectra_s3(ds3.DeleteBucketSpectraS3Request(bucket_name=bucket, force=True))
 
-    def test_put_and_get_all_objects_in_directory(self):
+    def test_put_and_get_all_objects_in_directory(self, prefix_dir: str = ""):
         bucket = f'ds3-python3-sdk-test-{uuid.uuid1()}'
         job_name = "python test job"
 
@@ -201,15 +201,20 @@ def test_put_and_get_all_objects_in_directory(self):
 
         # retrieve the objects from the BP
         destination = tempfile.TemporaryDirectory(prefix="ds3-python3-sdk-dst-")
+
+        # retrieve the objects in prefix directory
+        prefixed_objects = [p for p in put_objects if p.object_name.startswith(prefix_dir)]
+
         job_ids = helpers.get_all_files_in_bucket(destination_dir=destination.name,
                                                   bucket=bucket,
-                                                  objects_per_bp_job=10,
-                                                  job_name=job_name)
+                                                  objects_per_bp_job=10, job_name=job_name,
+                                                  prefix=prefix_dir)
 
-        self.assertGreaterEqual(len(job_ids), 2, "multiple job ids returned")
+        if len(prefix_dir) == 0:
+            self.assertGreaterEqual(len(job_ids), 2, "multiple job ids returned")
 
-        # verify all the files and directories were retrieved
-        for put_object in put_objects:
+        # verify all the files in prefix were returned
+        for put_object in prefixed_objects:
             obj_destination = os.path.join(destination.name,
                                            ds3Helpers.object_name_to_file_path(put_object.object_name))
             if put_object.object_name.endswith('/'):
@@ -228,6 +233,12 @@ def test_put_and_get_all_objects_in_directory(self):
         destination.cleanup()
         client.delete_bucket_spectra_s3(ds3.DeleteBucketSpectraS3Request(bucket_name=bucket, force=True))
 
+    def test_put_and_get_all_objects_in_directory_with_prefix(self):
+        self.test_put_and_get_all_objects_in_directory(prefix_dir="dir-0/sub-dir-1/")
+
+    def test_put_and_get_all_objects_in_directory_with_bad_prefix(self):
+        self.test_put_and_get_all_objects_in_directory(prefix_dir="not/gonna/match/")
+
     def test_put_all_objects_in_directory_with_md5_checksum(self):
         self.put_all_objects_in_directory_with_checksum(checksum_type='MD5')