#

# Licensed to the Apache Software Foundation (ASF) under one or more

# contributor license agreements. See the NOTICE file distributed with

# this work for additional information regarding copyright ownership.

# The ASF licenses this file to You under the Apache License, Version 2.0

# (the "License"); you may not use this file except in compliance with

# the License. You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License 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.

#

"""AWS S3 client

"""

# pytype: skip-file

import errno

import io

import logging

import re

import time

import traceback

from apache_beam.io.aws.clients.s3 import messages

from apache_beam.io.filesystemio import Downloader

from apache_beam.io.filesystemio import DownloaderStream

from apache_beam.io.filesystemio import Uploader

from apache_beam.io.filesystemio import UploaderStream

from apache_beam.utils import retry

try:

# pylint: disable=wrong-import-order, wrong-import-position

# pylint: disable=ungrouped-imports

from apache_beam.io.aws.clients.s3 import boto3_client

BOTO3_INSTALLED = True

except ImportError:

BOTO3_INSTALLED = False

MAX_BATCH_OPERATION_SIZE = 100

def parse_s3_path(s3_path, object_optional=False):

"""Return the bucket and object names of the given s3:// path."""

match = re.match('^s3://([^/]+)/(.*)$', s3_path)

if match is None or (match.group(2) == '' and not object_optional):

raise ValueError('S3 path must be in the form s3://<bucket>/<object>.')

return match.group(1), match.group(2)

class S3IO(object):

"""S3 I/O client."""

def __init__(self, client=None, options=None):

if client is None and options is None:

raise ValueError('Must provide one of client or options')

if client is not None:

self.client = client

elif BOTO3_INSTALLED:

self.client = boto3_client.Client(options=options)

else:

message = 'AWS dependencies are not installed, and no alternative ' \

'client was provided to S3IO.'

raise RuntimeError(message)

def open(

self,

filename,

mode='r',

read_buffer_size=16 * 1024 * 1024,

mime_type='application/octet-stream'):

"""Open an S3 file path for reading or writing.

Args:

filename (str): S3 file path in the form ``s3://<bucket>/<object>``.

mode (str): ``'r'`` for reading or ``'w'`` for writing.

read_buffer_size (int): Buffer size to use during read operations.

mime_type (str): Mime type to set for write operations.

Returns:

S3 file object.

Raises:

ValueError: Invalid open file mode.

"""

if mode == 'r' or mode == 'rb':

downloader = S3Downloader(

self.client, filename, buffer_size=read_buffer_size)

return io.BufferedReader(

DownloaderStream(downloader, mode=mode), buffer_size=read_buffer_size)

elif mode == 'w' or mode == 'wb':

uploader = S3Uploader(self.client, filename, mime_type)

return io.BufferedWriter(

UploaderStream(uploader, mode=mode), buffer_size=128 * 1024)

else:

raise ValueError('Invalid file open mode: %s.' % mode)

def list_files(self, path, with_metadata=False):

"""Lists files matching the prefix.

Args:

path: S3 file path pattern in the form s3://<bucket>/[name].

with_metadata: Experimental. Specify whether returns file metadata.

Returns:

If ``with_metadata`` is False: generator of tuple(file name, size); if

``with_metadata`` is True: generator of

tuple(file name, tuple(size, timestamp)).

"""

bucket, prefix = parse_s3_path(path, object_optional=True)

request = messages.ListRequest(bucket=bucket, prefix=prefix)

file_info = set()

counter = 0

start_time = time.time()

if with_metadata:

logging.debug("Starting the file information of the input")

else:

logging.debug("Starting the size estimation of the input")

while True:

#The list operation will raise an exception

#when trying to list a nonexistent S3 path.

#This should not be an issue here.

#Ignore this exception or it will break the procedure.

try:

response = retry.with_exponential_backoff(

retry_filter=retry.retry_on_server_errors_and_timeout_filter)(

self.client.list)(

request)

except messages.S3ClientError as e:

if e.code == 404:

break

else:

raise e

for item in response.items:

file_name = 's3://%s/%s' % (bucket, item.key)

if file_name not in file_info:

file_info.add(file_name)

counter += 1

if counter % 10000 == 0:

if with_metadata:

logging.info(

"Finished computing file information of: %s files",

len(file_info))

else:

logging.info(

"Finished computing size of: %s files", len(file_info))

if with_metadata:

yield file_name, (

item.size, self._updated_to_seconds(item.last_modified))

else:

yield file_name, item.size

if response.next_token:

request.continuation_token = response.next_token

else:

break

logging.log(

# do not spam logs when list_files is likely used to check empty folder

logging.INFO if counter > 0 else logging.DEBUG,

"Finished listing %s files in %s seconds.",

counter,

time.time() - start_time)

return file_info

def checksum(self, path):

"""Looks up the checksum of an S3 object.

Args:

path: S3 file path pattern in the form s3://<bucket>/<name>.

"""

return self._s3_object(path).etag

@retry.with_exponential_backoff(

retry_filter=retry.retry_on_server_errors_and_timeout_filter)

def copy(self, src, dest):

"""Copies a single S3 file object from src to dest.

Args:

src: S3 file path pattern in the form s3://<bucket>/<name>.

dest: S3 file path pattern in the form s3://<bucket>/<name>.

Raises:

TimeoutError: on timeout.

"""

src_bucket, src_key = parse_s3_path(src)

dest_bucket, dest_key = parse_s3_path(dest)

request = messages.CopyRequest(src_bucket, src_key, dest_bucket, dest_key)

self.client.copy(request)

# We intentionally do not decorate this method with a retry, since the

# underlying copy and delete operations are already idempotent operations

# protected by retry decorators.

def copy_paths(self, src_dest_pairs):

"""Copies the given S3 objects from src to dest. This can handle directory

or file paths.

Args:

src_dest_pairs: list of (src, dest) tuples of s3://<bucket>/<name> file

paths to copy from src to dest

Returns: List of tuples of (src, dest, exception) in the same order as the

src_dest_pairs argument, where exception is None if the operation

succeeded or the relevant exception if the operation failed.

"""

if not src_dest_pairs: return []

results = []

for src_path, dest_path in src_dest_pairs:

# Copy a directory with self.copy_tree

if src_path.endswith('/') and dest_path.endswith('/'):

try:

results += self.copy_tree(src_path, dest_path)

except messages.S3ClientError as err:

results.append((src_path, dest_path, err))

# Copy individual files with self.copy

elif not src_path.endswith('/') and not dest_path.endswith('/'):

src_bucket, src_key = parse_s3_path(src_path)

dest_bucket, dest_key = parse_s3_path(dest_path)

request = messages.CopyRequest(

src_bucket, src_key, dest_bucket, dest_key)

try:

self.client.copy(request)

results.append((src_path, dest_path, None))

except messages.S3ClientError as err:

results.append((src_path, dest_path, err))

# Mismatched paths (one directory, one non-directory) get an error result

else:

e = messages.S3ClientError(

"Can't copy mismatched paths (one directory, one non-directory):" +

' %s, %s' % (src_path, dest_path),

400)

results.append((src_path, dest_path, e))

return results

# We intentionally do not decorate this method with a retry, since the

# underlying copy and delete operations are already idempotent operations

# protected by retry decorators.

def copy_tree(self, src, dest):

"""Renames the given S3 directory and it's contents recursively

from src to dest.

Args:

src: S3 file path pattern in the form s3://<bucket>/<name>/.

dest: S3 file path pattern in the form s3://<bucket>/<name>/.

Returns:

List of tuples of (src, dest, exception) where exception is None if the

operation succeeded or the relevant exception if the operation failed.

"""

assert src.endswith('/')

assert dest.endswith('/')

results = []

for entry, _ in self.list_files(src):

rel_path = entry[len(src):]

try:

self.copy(entry, dest + rel_path)

results.append((entry, dest + rel_path, None))

except messages.S3ClientError as e:

results.append((entry, dest + rel_path, e))

return results

@retry.with_exponential_backoff(

retry_filter=retry.retry_on_server_errors_and_timeout_filter)

def delete(self, path):

"""Deletes a single S3 file object from src to dest.

Args:

src: S3 file path pattern in the form s3://<bucket>/<name>/.

dest: S3 file path pattern in the form s3://<bucket>/<name>/.

Returns:

List of tuples of (src, dest, exception) in the same order as the

src_dest_pairs argument, where exception is None if the operation

succeeded or the relevant exception if the operation failed.

"""

bucket, object_path = parse_s3_path(path)

request = messages.DeleteRequest(bucket, object_path)

try:

self.client.delete(request)

except messages.S3ClientError as e:

if e.code == 404:

return # Same behavior as GCS - don't surface a 404 error

else:

logging.error('HTTP error while deleting file %s: %s', path, 3)

raise e

# We intentionally do not decorate this method with a retry, since the

# underlying copy and delete operations are already idempotent operations

# protected by retry decorators.

def delete_paths(self, paths):

"""Deletes the given S3 objects from src to dest. This can handle directory

or file paths.

Args:

src: S3 file path pattern in the form s3://<bucket>/<name>/.

dest: S3 file path pattern in the form s3://<bucket>/<name>/.

Returns:

List of tuples of (src, dest, exception) in the same order as the

src_dest_pairs argument, where exception is None if the operation

succeeded or the relevant exception if the operation failed.

"""

directories, not_directories = [], []

for path in paths:

if path.endswith('/'): directories.append(path)

else: not_directories.append(path)

results = {}

for directory in directories:

dir_result = dict(self.delete_tree(directory))

results.update(dir_result)

not_directory_results = dict(self.delete_files(not_directories))

results.update(not_directory_results)

return results

# We intentionally do not decorate this method with a retry, since the

# underlying copy and delete operations are already idempotent operations

# protected by retry decorators.

def delete_files(self, paths, max_batch_size=1000):

"""Deletes the given S3 file object from src to dest.

Args:

paths: List of S3 file paths in the form s3://<bucket>/<name>

max_batch_size: Largest number of keys to send to the client to be deleted

simultaneously

Returns: List of tuples of (path, exception) in the same order as the paths

argument, where exception is None if the operation succeeded or

the relevant exception if the operation failed.

"""

if not paths: return []

# Sort paths into bucket: [keys]

buckets, keys = zip(*[parse_s3_path(path) for path in paths])

grouped_keys = {bucket: [] for bucket in buckets}

for bucket, key in zip(buckets, keys):

grouped_keys[bucket].append(key)

# For each bucket, delete minibatches of keys

results = {}

for bucket, keys in grouped_keys.items():

for i in range(0, len(keys), max_batch_size):

minibatch_keys = keys[i:i + max_batch_size]

results.update(self._delete_minibatch(bucket, minibatch_keys))

# Organize final results

final_results = [(path, results[parse_s3_path(path)]) for path in paths]

return final_results

@retry.with_exponential_backoff(

retry_filter=retry.retry_on_server_errors_and_timeout_filter)

def _delete_minibatch(self, bucket, keys):

"""A helper method. Boto3 allows batch deletions

for files within the same bucket.

Args:

bucket: String bucket name

keys: List of keys to be deleted in the bucket

Returns: dict of the form {(bucket, key): error}, where error is None if the

operation succeeded

"""

request = messages.DeleteBatchRequest(bucket, keys)

results = {}

try:

response = self.client.delete_batch(request)

for key in response.deleted:

results[(bucket, key)] = None

for key, error in zip(response.failed, response.errors):

results[(bucket, key)] = error

except messages.S3ClientError as e:

for key in keys:

results[(bucket, key)] = e

return results

# We intentionally do not decorate this method with a retry, since the

# underlying copy and delete operations are already idempotent operations

# protected by retry decorators.

def delete_tree(self, root):

"""Deletes all objects under the given S3 directory.

Args:

path: S3 root path in the form s3://<bucket>/<name>/ (ending with a "/")

Returns: List of tuples of (path, exception), where each path is an object

under the given root. exception is None if the operation succeeded

or the relevant exception if the operation failed.

"""

assert root.endswith('/')

paths = [p for p, _ in self.list_files(root)]

return self.delete_files(paths)

def size(self, path):

"""Returns the size of a single S3 object.

This method does not perform glob expansion. Hence the given path must be

for a single S3 object.

Returns: size of the S3 object in bytes.

"""

return self._s3_object(path).size

# We intentionally do not decorate this method with a retry, since the

# underlying copy and delete operations are already idempotent operations

# protected by retry decorators.

def rename(self, src, dest):

"""Renames the given S3 object from src to dest.

Args:

src: S3 file path pattern in the form s3://<bucket>/<name>.

dest: S3 file path pattern in the form s3://<bucket>/<name>.

"""

self.copy(src, dest)

self.delete(src)

def last_updated(self, path):

"""Returns the last updated epoch time of a single S3 object.

This method does not perform glob expansion. Hence the given path must be

for a single S3 object.

Returns: last updated time of the S3 object in second.

"""

return self._updated_to_seconds(self._s3_object(path).last_modified)

def exists(self, path):

"""Returns whether the given S3 object exists.

Args:

path: S3 file path pattern in the form s3://<bucket>/<name>.

"""

try:

self._s3_object(path)

return True

except messages.S3ClientError as e:

if e.code == 404:

# HTTP 404 indicates that the file did not exist

return False

else:

# We re-raise all other exceptions

raise

def _status(self, path):

"""For internal use only; no backwards-compatibility guarantees.

Returns supported fields (checksum, last_updated, size) of a single object

as a dict at once.

This method does not perform glob expansion. Hence the given path must be

for a single S3 object.

Returns: dict of fields of the S3 object.

"""

s3_object = self._s3_object(path)

file_status = {}

if hasattr(s3_object, 'etag'):

file_status['checksum'] = s3_object.etag

if hasattr(s3_object, 'last_modified'):

file_status['last_updated'] = self._updated_to_seconds(

s3_object.last_modified)

if hasattr(s3_object, 'size'):

file_status['size'] = s3_object.size

return file_status

@retry.with_exponential_backoff(

retry_filter=retry.retry_on_server_errors_and_timeout_filter)

def _s3_object(self, path):

"""Returns a S3 object metadata for the given path

This method does not perform glob expansion. Hence the given path must be

for a single S3 object.

Returns: S3 object metadata.

"""

bucket, object = parse_s3_path(path)

request = messages.GetRequest(bucket, object)

return self.client.get_object_metadata(request)

@staticmethod

def _updated_to_seconds(updated):

"""Helper function transform the updated field of response to seconds"""

return (

time.mktime(updated.timetuple()) - time.timezone +

updated.microsecond / 1000000.0)

def rename_files(self, src_dest_pairs):

"""Renames the given S3 objects from src to dest.

Args:

src_dest_pairs: list of (src, dest) tuples of s3://<bucket>/<name> file

paths to rename from src to dest

Returns: List of tuples of (src, dest, exception) in the same order as the

src_dest_pairs argument, where exception is None if the operation

succeeded or the relevant exception if the operation failed.

"""

if not src_dest_pairs: return []

# TODO: Throw value error if path has directory

for src, dest in src_dest_pairs:

if src.endswith('/') or dest.endswith('/'):

raise ValueError('Cannot rename a directory')

copy_results = self.copy_paths(src_dest_pairs)

paths_to_delete = [src for (src, _, err) in copy_results if err is None]

delete_results = self.delete_files(paths_to_delete)

delete_results_dict = {src: err for (src, err) in delete_results}

rename_results = []

for src, dest, err in copy_results:

if err is not None: rename_results.append((src, dest, err))

elif delete_results_dict[src] is not None:

rename_results.append((src, dest, delete_results_dict[src]))

else:

rename_results.append((src, dest, None))

return rename_results

class S3Downloader(Downloader):

def __init__(self, client, path, buffer_size):

self._client = client

self._path = path

self._bucket, self._name = parse_s3_path(path)

self._buffer_size = buffer_size

# Get object state.

self._get_request = (

messages.GetRequest(bucket=self._bucket, object=self._name))

try:

metadata = self._get_object_metadata(self._get_request)

except messages.S3ClientError as e:

if e.code == 404:

raise IOError(errno.ENOENT, 'Not found: %s' % self._path)

else:

logging.error('HTTP error while requesting file %s: %s', self._path, 3)

raise

self._size = metadata.size

@retry.with_exponential_backoff(

retry_filter=retry.retry_on_server_errors_and_timeout_filter)

def _get_object_metadata(self, get_request):

return self._client.get_object_metadata(get_request)

@property

def size(self):

return self._size

def get_range(self, start, end):

return self._client.get_range(self._get_request, start, end)

class S3Uploader(Uploader):

def __init__(self, client, path, mime_type='application/octet-stream'):

self._client = client

self._path = path

self._bucket, self._name = parse_s3_path(path)

self._mime_type = mime_type

self.part_number = 1

self.buffer = b''

self.last_error = None

self.upload_id = None

self.parts = []

self._start_upload()

# There is retry logic in the underlying transfer library but we should make

# it more explicit so we can control the retry parameters.

@retry.no_retries # Using no_retries marks this as an integration point.

def _start_upload(self):

# The uploader by default transfers data in chunks of 1024 * 1024 bytes at

# a time, buffering writes until that size is reached.

try:

request = messages.UploadRequest(

self._bucket, self._name, self._mime_type)

response = self._client.create_multipart_upload(request)

self.upload_id = response.upload_id

except Exception as e: # pylint: disable=broad-except

logging.error(

'Error in _start_upload while inserting file %s: %s',

self._path,

traceback.format_exc())

self.last_error = e

raise e

def put(self, data):

MIN_WRITE_SIZE = 5 * 1024 * 1024

MAX_WRITE_SIZE = 5 * 1024 * 1024 * 1024

# TODO: Byte strings might not be the most performant way to handle this

self.buffer += data.tobytes()

while len(self.buffer) >= MIN_WRITE_SIZE:

# Take the first chunk off the buffer and write it to S3

chunk = self.buffer[:MAX_WRITE_SIZE]

self._write_to_s3(chunk)

# Remove the written chunk from the buffer

self.buffer = self.buffer[MAX_WRITE_SIZE:]

def _write_to_s3(self, data):

try:

request = messages.UploadPartRequest(

self._bucket, self._name, self.upload_id, self.part_number, data)

response = self._client.upload_part(request)

self.parts.append({

'ETag': response.etag, 'PartNumber': response.part_number

})

self.part_number = self.part_number + 1

except messages.S3ClientError as e:

self.last_error = e

if e.code == 404:

raise IOError(errno.ENOENT, 'Not found: %s' % self._path)

else:

logging.error('HTTP error while requesting file %s: %s', self._path, 3)

raise

def finish(self):

if len(self.buffer) > 0:

# writes with zero length or mid size files between

# MIN_WRITE_SIZE = 5 * 1024 * 1024

# MAX_WRITE_SIZE = 5 * 1024 * 1024 * 1024

# as we will reach this finish() with len(self.buffer) == 0

# which will fail

self._write_to_s3(self.buffer)

if self.last_error is not None:

raise self.last_error # pylint: disable=raising-bad-type

request = messages.CompleteMultipartUploadRequest(

self._bucket, self._name, self.upload_id, self.parts)

self._client.complete_multipart_upload(request)