feature: add S3Downloader.list(s3_uri) functionality (#283)

Adds the ability to call S3Downloader.list(base_s3_uri) on a base S3 path, which returns a list of S3 URIs corresponding to valid S3 objects under that path.

Author: knakad

src/sagemaker/s3.py

@@ -13,6 +13,8 @@
 """This module contains Enums and helper methods related to S3."""
 from __future__ import print_function, absolute_import
 
+import os
+
 from six.moves.urllib.parse import urlparse
 from sagemaker.session import Session
 
@@ -42,7 +44,7 @@ def upload(local_path, desired_s3_uri, kms_key=None, session=None):
         Args:
             local_path (str): A local path to a file or directory.
             desired_s3_uri (str): The desired S3 uri to upload to.
-            kms_key (str): A KMS key to be provided as an extra argument.
+            kms_key (str): The KMS key to use to encrypt the files.
             session (sagemaker.session.Session): Session object which
                 manages interactions with Amazon SageMaker APIs and any other
                 AWS services needed. If not specified, the estimator creates one
@@ -70,7 +72,7 @@ def upload_string_as_file_body(body, desired_s3_uri=None, kms_key=None, session=
         Args:
             body (str): String representing the body of the file.
             desired_s3_uri (str): The desired S3 uri to upload to.
-            kms_key (str): A KMS key to be provided as an extra argument.
+            kms_key (str): The KMS key to use to encrypt the files.
             session (sagemaker.session.Session): AWS session to use. Automatically
                 generates one if not provided.
 
@@ -98,7 +100,7 @@ def download(s3_uri, local_path, kms_key=None, session=None):
         Args:
             s3_uri (str): An S3 uri to download from.
             local_path (str): A local path to download the file(s) to.
-            kms_key (str): A KMS key to be provided as an extra argument.
+            kms_key (str): The KMS key to use to decrypt the files.
             session (sagemaker.session.Session): Session object which
                 manages interactions with Amazon SageMaker APIs and any other
                 AWS services needed. If not specified, the estimator creates one
@@ -133,3 +135,22 @@ def read_file(s3_uri, session=None):
         bucket, key_prefix = parse_s3_url(url=s3_uri)
 
         return sagemaker_session.read_s3_file(bucket=bucket, key_prefix=key_prefix)
+
+    @staticmethod
+    def list(s3_uri, session=None):
+        """Static method that lists the contents of an S3 uri.
+
+        Args:
+            s3_uri (str): The S3 base uri to list objects in.
+            session (sagemaker.session.Session): AWS session to use. Automatically
+                generates one if not provided.
+
+        Returns:
+            [str]: The list of S3 URIs in the given S3 base uri.
+
+        """
+        sagemaker_session = session or Session()
+        bucket, key_prefix = parse_s3_url(url=s3_uri)
+
+        file_keys = sagemaker_session.list_s3_files(bucket=bucket, key_prefix=key_prefix)
+        return [os.path.join("s3://", bucket, file_key) for file_key in file_keys]

src/sagemaker/session.py

@@ -203,7 +203,7 @@ def upload_string_as_file_body(self, body, bucket, key, kms_key=None):
                 default bucket of the ``Session`` is used (if default bucket does not exist, the
                 ``Session`` creates it).
             key (str): S3 object key. This is the s3 path to the file.
-            kms_key (str): The KMS key to use for decrypting the file.
+            kms_key (str): The KMS key to use for encrypting the file.
 
         Returns:
             str: The S3 URI of the uploaded file(s). If a file is specified in the path argument,
@@ -291,7 +291,24 @@ def read_s3_file(self, bucket, key_prefix):
         # Explicitly passing a None kms_key to boto3 throws a validation error.
         s3_object = s3.get_object(Bucket=bucket, Key=key_prefix)
 
-        return s3_object["Body"].read()
+        return s3_object["Body"].read().decode("utf-8")
+
+    def list_s3_files(self, bucket, key_prefix):
+        """Lists the S3 files given an S3 bucket and key.
+
+        Args:
+            bucket (str): Name of the S3 Bucket to download from.
+            key_prefix (str): S3 object key name prefix.
+
+        Returns:
+            [str]: The list of files at the S3 path.
+
+        """
+        s3 = self.boto_session.resource("s3")
+
+        s3_bucket = s3.Bucket(name=bucket)
+        s3_objects = s3_bucket.objects.filter(Prefix=key_prefix).all()
+        return [s3_object.key for s3_object in s3_objects]
 
     def default_bucket(self):
         """Return the name of the default bucket to use in relevant Amazon SageMaker interactions.

tests/integ/test_s3.py

@@ -0,0 +1,67 @@
+# Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import os
+import uuid
+
+import pytest
+
+from sagemaker.s3 import S3Uploader
+from sagemaker.s3 import S3Downloader
+
+from tests.integ.kms_utils import get_or_create_kms_key
+
+
+@pytest.fixture(scope="module")
+def s3_files_kms_key(sagemaker_session):
+    return get_or_create_kms_key(sagemaker_session=sagemaker_session)
+
+
+def test_statistics_object_creation_from_s3_uri_with_customizations(
+    sagemaker_session, s3_files_kms_key
+):
+    file_1_body = "First File Body."
+    file_1_name = "first_file.txt"
+    file_2_body = "Second File Body."
+    file_2_name = "second_file.txt"
+
+    my_uuid = str(uuid.uuid4())
+
+    base_s3_uri = os.path.join(
+        "s3://", sagemaker_session.default_bucket(), "integ-test-test-s3-list", my_uuid
+    )
+    file_1_s3_uri = os.path.join(base_s3_uri, file_1_name)
+    file_2_s3_uri = os.path.join(base_s3_uri, file_2_name)
+
+    S3Uploader.upload_string_as_file_body(
+        body=file_1_body,
+        desired_s3_uri=file_1_s3_uri,
+        kms_key=s3_files_kms_key,
+        session=sagemaker_session,
+    )
+
+    S3Uploader.upload_string_as_file_body(
+        body=file_2_body,
+        desired_s3_uri=file_2_s3_uri,
+        kms_key=s3_files_kms_key,
+        session=sagemaker_session,
+    )
+
+    s3_uris = S3Downloader.list(s3_uri=base_s3_uri, session=sagemaker_session)
+
+    assert file_1_name in s3_uris[0]
+    assert file_2_name in s3_uris[1]
+
+    assert file_1_body == S3Downloader.read_file(s3_uri=s3_uris[0], session=sagemaker_session)
+    assert file_2_body == S3Downloader.read_file(s3_uri=s3_uris[1], session=sagemaker_session)