dvc/dvc/command/data_sync.py

import base64
import hashlib
import os

import fasteners
from boto.s3.connection import S3Connection
from google.cloud import storage as gc

from dvc.command.base import CmdBase
from dvc.command.init import CmdInit
from dvc.logger import Logger
from dvc.exceptions import DvcException
from dvc.runtime import Runtime
from dvc.system import System


class DataSyncError(DvcException):
    def __init__(self, msg):
        DvcException.__init__(self, 'Data sync error: {}'.format(msg))


def sizeof_fmt(num, suffix='B'):
    for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:
        if abs(num) < 1024.0:
            return "%3.1f%s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f%s%s" % (num, 'Y', suffix)


def percent_cb(complete, total):
    Logger.debug('{} transferred out of {}'.format(sizeof_fmt(complete), sizeof_fmt(total)))


def file_md5(fname):
    """ get the (md5 hexdigest, md5 digest) of a file """
    hash_md5 = hashlib.md5()
    with open(fname, "rb") as f:
        for chunk in iter(lambda: f.read(1024*1000), b""):
            hash_md5.update(chunk)
    return (hash_md5.hexdigest(), hash_md5.digest())


class CmdDataSync(CmdBase):
    def __init__(self, settings):
        super(CmdDataSync, self).__init__(settings)

    def define_args(self, parser):
        parser.add_argument('targets',
                            metavar='',
                            help='File or directory to sync.',
                            nargs='*')

    def run(self):
        if self.is_locker:
            lock = fasteners.InterProcessLock(self.git.lock_file)
            gotten = lock.acquire(timeout=5)
            if not gotten:
                Logger.info('Cannot perform the cmd since DVC is busy and locked. Please retry the cmd later.')
                return 1

        good = self.config.sanity_check()

        if not good[0]:
            Logger.error('config \'%s\' is not correctly setup.  Please fix:' % Runtime.CONFIG)
            for e in good[1]:
                Logger.error('    ' + e)
            return 1

        try:
            for target in self.parsed_args.targets:
                data_item = self.settings.path_factory.data_item(target)
                if System.islink(target):
                    self.sync_symlink(data_item)
                elif os.path.isdir(target):
                    self.sync_dir(target)
                else:
                    raise DataSyncError('File "{}" does not exit'.format(target))
        finally:
            if self.is_locker:
                lock.release()

    def sync_dir(self, dir):
        for file in os.listdir(dir):
            try:
                fname = None
                fname = os.path.join(dir, file)
                if os.path.isdir(fname):
                    self.sync_dir(fname)
                elif System.islink(fname):
                    self.sync_symlink(self.settings.path_factory.data_item(fname))
                else:
                    raise DvcException('Unsupported file type "{}"'.format(fname))
            except DataSyncError as ex:
                Logger.debug(ex)
            except Exception as ex:
                Logger.error('Cannot sync file {}: {}'.format(fname, ex))
        pass

    def sync_symlink(self, data_item):
        if os.path.isfile(data_item.resolved_cache.dvc):
            self.sync_to_cloud(data_item)
        else:
            self.sync_from_cloud(data_item)

    def _get_bucket_aws(self):
        """ get a bucket object, aws """
        conn = S3Connection(self.config.aws_access_key_id, self.config.aws_secret_access_key, host=self.config.aws_region_host)
        bucket_name = self.config.storage_bucket
        bucket = conn.lookup(bucket_name)
        if bucket is None:
            raise DataSyncError('Storage path is not setup correctly')
        return bucket

    def _sync_from_cloud_aws(self, item):
        """ sync from cloud, aws version """

        bucket = self._get_bucket_aws()

        key_name = self.cache_file_key(item.resolved_cache.dvc)
        key = bucket.get_key(key_name)
        if not key:
            raise DataSyncError('File "{}" does not exist in the cloud'.format(key_name))

        Logger.info('Downloading cache file from S3 "{}/{}"'.format(bucket.name, key_name))
        key.get_contents_to_filename(item.resolved_cache.relative, cb=percent_cb)
        Logger.info('Downloading completed')

    def _sync_to_cloud_aws(self, data_item):
        """ sync_to_cloud, aws version """

        aws_key = self.cache_file_key(data_item.resolved_cache.dvc)
        bucket = self._get_bucket_aws()
        key = bucket.get_key(aws_key)
        if key:
            Logger.debug('File already uploaded to the cloud. Checksum validation...')

            md5_cloud = key.etag[1:-1]
            md5_local = file_md5(data_item.resolved_cache.dvc)[0]
            if md5_cloud == md5_local:
                Logger.debug('File checksum matches. No uploading is needed.')
                return

            Logger.debug('Checksum miss-match. Re-uploading is required.')

        Logger.info('Uploading cache file "{}" to S3 "{}"'.format(data_item.resolved_cache.relative, aws_key))
        key = bucket.new_key(aws_key)
        key.set_contents_from_filename(data_item.resolved_cache.relative, cb=percent_cb)
        Logger.info('Uploading completed')


    def _get_bucket_gc(self):
        """ get a bucket object, gc """
        client = gc.Client(project=self.config.gc_project_name)
        bucket = client.bucket(self.config.storage_bucket)
        if not bucket.exists():
            raise DataSyncError('sync up: google cloud bucket {} doesn\'t exist'.format(self.config.storage_bucket))
        return bucket

    def _sync_from_cloud_gcp(self, item):
        """ sync from cloud, gcp version
        Params:
            TODO
        """

        bucket = self._get_bucket_gc()
        key = self.cache_file_key(item.resolved_cache.dvc)

        blob = bucket.get_blob(key)
        if not blob:
            raise DataSyncError('File "{}" does not exist in the cloud'.format(key))

        Logger.info('Downloading cache file from gc "{}/{}"'.format(bucket.name, key))

        blob.download_to_filename(item.resolved_cache.dvc)
        Logger.info('Downloading completed')

    def _sync_to_cloud_gcp(self, data_item):
        """ sync_to_cloud, gcp version """

        bucket = self._get_bucket_gc()
        blob_name = self.cache_file_key(data_item.resolved_cache.dvc)

        blob = bucket.get_blob(blob_name)
        if blob is not None and blob.exists():
            b64_encoded_md5 = base64.b64encode(file_md5(data_item.resolved_cache.dvc)[1])

            if blob.md5_hash == b64_encoded_md5:
                Logger.debug('checksum %s matches.  Skipping upload' % data_item.cache.relative)
                return
            Logger.debug('checksum %s mismatch.  re-uploading' % data_item.cache.relative)

        Logger.info('uploading cache file "{} to gc "{}"'.format(data_item.cache.relative, blob_name))

        blob = bucket.blob(blob_name)
        blob.upload_from_filename(data_item.resolved_cache.relative)
        Logger.info('uploading %s completed' % data_item.resolved_cache.relative)

    def sync_from_cloud(self, item):
        cloud = self.settings.config.cloud
        assert cloud in ['aws', 'gcp'], 'unknown cloud %s' % cloud
        if cloud == 'aws':
            return self._sync_from_cloud_aws(item)
        elif cloud == 'gcp':
            return self._sync_from_cloud_gcp(item)

    def sync_to_cloud(self, data_item):
        cloud = self.settings.config.cloud
        assert cloud in ['aws', 'gcp'], 'unknown cloud %s' % cloud
        if cloud == 'aws':
            return self._sync_to_cloud_aws(data_item)
        elif cloud == 'gcp':
            return self._sync_to_cloud_gcp(data_item)


if __name__ == '__main__':
    Runtime.run(CmdDataSync)