blender-asset-tracer/blender_asset_tracer/pack/s3.py

"""Amazon S3-compatible uploader."""
import typing

import hashlib
import logging
import pathlib
import threading

from . import Packer, transfer

log = logging.getLogger(__name__)


def compute_md5(filepath: pathlib.Path) -> str:
    hasher = hashlib.md5()
    with filepath.open('rb') as infile:
        while True:
            block = infile.read(10240)
            if not block:
                break
            hasher.update(block)
    md5 = hasher.hexdigest()
    return md5


class S3Packer(Packer):
    """Creates BAT Packs on S3-compatible storage."""

    def __init__(self, *args, endpoint, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        import boto3

        # Create a session so that credentials can be read from the [endpoint]
        # section in ~/.aws/credentials.
        # See https://boto3.readthedocs.io/en/latest/guide/configuration.html#guide-configuration
        self.session = boto3.Session(profile_name=endpoint)
        self.client = self.session.client('s3', endpoint_url='https://%s' % endpoint)

    def set_credentials(self,
                        endpoint: str,
                        access_key_id: str,
                        secret_access_key: str):
        """Set S3 credentials."""
        import boto3

        self.client = boto3.client('s3',
                                   endpoint_url=endpoint,
                                   aws_access_key_id=access_key_id,
                                   aws_secret_access_key=secret_access_key)

    def _create_file_transferer(self) -> transfer.FileTransferer:
        return S3Transferrer(self.client)


class S3Transferrer(transfer.FileTransferer):
    """Copies or moves files in source directory order."""

    class AbortUpload(Exception):
        """Raised from the upload callback to abort an upload."""

    def __init__(self, botoclient) -> None:
        super().__init__()
        self.client = botoclient

    def run(self) -> None:
        files_transferred = 0
        files_skipped = 0

        for src, dst, act in self.iter_queue():
            try:
                did_upload = self.upload_file(src, dst)
                files_transferred += did_upload
                files_skipped += not did_upload

                if act == transfer.Action.MOVE:
                    self.delete_file(src)
            except Exception:
                # We have to catch exceptions in a broad way, as this is running in
                # a separate thread, and exceptions won't otherwise be seen.
                log.exception('Error transferring %s to %s', src, dst)
                # Put the files to copy back into the queue, and abort. This allows
                # the main thread to inspect the queue and see which files were not
                # copied. The one we just failed (due to this exception) should also
                # be reported there.
                self.queue.put((src, dst, act))
                return

        if files_transferred:
            log.info('Transferred %d files', files_transferred)
        if files_skipped:
            log.info('Skipped %d files', files_skipped)

    def upload_file(self, src: pathlib.Path, dst: pathlib.Path) -> bool:
        """Upload a file to an S3 bucket.

        The first part of 'dst' is used as the bucket name, the remained as the
        path inside the bucket.

        :returns: True if the file was uploaded, False if it was skipped.
        """
        bucket = dst.parts[0]
        dst_path = pathlib.Path(*dst.parts[1:])
        md5 = compute_md5(src)
        key = str(dst_path)

        existing_md5, existing_size = self.get_metadata(bucket, key)
        if md5 == existing_md5 and src.stat().st_size == existing_size:
            log.debug('skipping %s, it already exists on the server with MD5 %s',
                      src, existing_md5)
            return False

        # TODO(Sybren): when queueing files inspect their size, and have a
        # callback that reports the total progress.
        log.info('Uploading %s', src)
        try:
            self.client.upload_file(str(src),
                                    Bucket=bucket,
                                    Key=key,
                                    Callback=self._upload_callback,
                                    ExtraArgs={'Metadata': {'md5': md5}})
        except self.AbortUpload:
            return False
        return True

    def _upload_callback(self, bytes_uploaded: int):
        if self.abort.is_set():
            log.warning('Interrupting ongoing upload')
            raise self.AbortUpload('interrupting ongoing upload')

    def delete_file(self, path: pathlib.Path):
        """Deletes a file, only logging a warning if deletion fails."""
        log.debug('Deleting %s, file has been uploaded', path)
        try:
            path.unlink()
        except IOError as ex:
            log.warning('Unable to delete %s: %s', path, ex)

    def get_metadata(self, bucket: str, key: str) -> typing.Tuple[str, int]:
        """Get MD5 sum and size on S3.

        :returns: the MD5 hexadecimal hash and the file size in bytes.
            If the file does not exist or has no known MD5 sum,
            returns ('', -1)
        """
        import botocore.exceptions

        try:
            info = self.client.head_object(Bucket=bucket, Key=key)
        except botocore.exceptions.ClientError as ex:
            error_code = ex.response.get('Error').get('Code', 'Unknown')
            # error_code already is a string, but this makes the code forward
            # compatible with a time where they use integer codes.
            if str(error_code) == '404':
                return '', -1
            raise ValueError('error response:' % ex.response) from None

        try:
            return info['Metadata']['md5'], info['ContentLength']
        except KeyError:
            return '', -1