Source code for apache_beam.io.aws.s3io

#
# 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


[docs]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)
[docs]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)
[docs] 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)
[docs] @retry.with_exponential_backoff( retry_filter=retry.retry_on_server_errors_and_timeout_filter) def list_prefix(self, path): """Lists files matching the prefix. Args: path: S3 file path pattern in the form s3://<bucket>/[name]. Returns: Dictionary of file name -> size. """ bucket, prefix = parse_s3_path(path, object_optional=True) request = messages.ListRequest(bucket=bucket, prefix=prefix) file_sizes = {} counter = 0 start_time = time.time() logging.info("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 = 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) file_sizes[file_name] = item.size counter += 1 if counter % 10000 == 0: logging.info("Finished computing size of: %s files", len(file_sizes)) if response.next_token: request.continuation_token = response.next_token else: break logging.info( "Finished listing %s files in %s seconds.", counter, time.time() - start_time) return file_sizes
[docs] @retry.with_exponential_backoff( retry_filter=retry.retry_on_server_errors_and_timeout_filter) def checksum(self, path): """Looks up the checksum of an S3 object. Args: path: S3 file path pattern in the form s3://<bucket>/<name>. """ bucket, object_path = parse_s3_path(path) request = messages.GetRequest(bucket, object_path) item = self.client.get_object_metadata(request) return item.etag
[docs] @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.
[docs] 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.
[docs] 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_prefix(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
[docs] @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.
[docs] 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.
[docs] 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.
[docs] 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 = self.list_prefix(root) return self.delete_files(paths)
[docs] @retry.with_exponential_backoff( retry_filter=retry.retry_on_server_errors_and_timeout_filter) 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. """ bucket, object_path = parse_s3_path(path) request = messages.GetRequest(bucket, object_path) item = self.client.get_object_metadata(request) return item.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.
[docs] 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)
[docs] @retry.with_exponential_backoff( retry_filter=retry.retry_on_server_errors_and_timeout_filter) 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. """ bucket, object = parse_s3_path(path) request = messages.GetRequest(bucket, object) datetime = self.client.get_object_metadata(request).last_modified return ( time.mktime(datetime.timetuple()) - time.timezone + datetime.microsecond / 1000000.0)
[docs] def exists(self, path): """Returns whether the given S3 object exists. Args: path: S3 file path pattern in the form s3://<bucket>/<name>. """ bucket, object = parse_s3_path(path) request = messages.GetRequest(bucket, object) try: self.client.get_object_metadata(request) 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
[docs] 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
[docs]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
[docs] def get_range(self, start, end): return self._client.get_range(self._get_request, start, end)
[docs]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
[docs] 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
[docs] def finish(self): 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)