diff --git a/dvc/command/base.py b/dvc/command/base.py index c58a72b17..722c77492 100644 --- a/dvc/command/base.py +++ b/dvc/command/base.py @@ -68,9 +68,6 @@ class CmdBase(object): def config(self): return self._settings.config - def cache_file_key(self, file): - return '{}/{}'.format(self.config.storage_prefix, file).strip('/') - @property def dvc_home(self): return self._settings.dvc_home @@ -147,12 +144,6 @@ class CmdBase(object): def run(self): pass - def get_cache_file_s3_name(self, cache_file): - cache_prefix_file_name = os.path.relpath(System.realpath(cache_file), System.realpath(self.git.git_dir)) - file_name = os.path.relpath(cache_prefix_file_name, self.config.cache_dir) - dvc_file_path_trim = file_name.replace(os.sep, '/').strip('/') - return self.config.aws_storage_prefix + '/' + dvc_file_path_trim - @staticmethod def warning_dvc_is_busy(): Logger.warn('Cannot perform the cmd since DVC is busy and locked. Please retry the cmd later.') diff --git a/dvc/command/data_sync.py b/dvc/command/data_sync.py index cc01009f1..32565d1a7 100644 --- a/dvc/command/data_sync.py +++ b/dvc/command/data_sync.py @@ -2,6 +2,8 @@ import base64 import hashlib import os +from multiprocessing.pool import ThreadPool + from boto.s3.connection import S3Connection from google.cloud import storage as gc @@ -10,34 +12,12 @@ from dvc.logger import Logger from dvc.exceptions import DvcException from dvc.runtime import Runtime from dvc.system import System - +from dvc.data_cloud import DataCloud class DataSyncError(DvcException): def __init__(self, msg): super(DataSyncError, self).__init__('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) @@ -50,158 +30,24 @@ class CmdDataSync(CmdBase): def run(self): with DvcLock(self.is_locker, self.git): - self.sync_sanity_check() + cloud = DataCloud(self.settings) + #FIXME add cmdline argument and config option for POOL_SIZE + POOL_SIZE = 10 + pool = ThreadPool(processes=POOL_SIZE) + targets = [] 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) + targets.append(target) elif os.path.isdir(target): - self.sync_dir(target) + for root, dirs, files in os.walk(target): + for f in files: + targets.append(os.path.join(root, f)) else: - raise DataSyncError('File "{}" does not exit'.format(target)) + raise DataSyncError('File "{}" does not exit'.format(target)) + + pool.map(cloud.sync, targets) pass - def sync_sanity_check(self): - good = self.config.sanity_check() - if not good[0]: - error = 'config \'%s\' is not correctly setup. Please fix:\n' % Runtime.CONFIG - for e in good[1]: - error += ' \n' + e - raise DataSyncError(error) - - def sync_dir(self, dir): - for file in os.listdir(dir): - try: - fname = None - fname = os.path.join(dir, file) - self.sync_object(fname) - except DataSyncError as ex: - Logger.debug(ex) - except Exception as ex: - Logger.error('Cannot sync file {}: {}'.format(fname, ex)) - pass - - def sync_object(self, fname): - 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)) - - 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 """ - - 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) diff --git a/dvc/command/import_file.py b/dvc/command/import_file.py index d031f0d8c..9defd2261 100644 --- a/dvc/command/import_file.py +++ b/dvc/command/import_file.py @@ -4,7 +4,7 @@ import re import requests from dvc.command.base import CmdBase, DvcLock -from dvc.command.data_sync import sizeof_fmt +from dvc.data_cloud import sizeof_fmt from dvc.logger import Logger from dvc.exceptions import DvcException from dvc.runtime import Runtime diff --git a/dvc/command/remove.py b/dvc/command/remove.py index 2b23bb926..34b8d06f5 100644 --- a/dvc/command/remove.py +++ b/dvc/command/remove.py @@ -6,7 +6,7 @@ from dvc.command.base import CmdBase, DvcLock from dvc.exceptions import DvcException from dvc.logger import Logger from dvc.runtime import Runtime - +from dvc.data_cloud import DataCloud class DataRemoveError(DvcException): def __init__(self, msg): @@ -30,6 +30,8 @@ class CmdDataRemove(CmdBase): def run(self): with DvcLock(self.is_locker, self.git): + self.cloud = DataCloud(self.settings) + if not self.remove_all_targets(): return 1 @@ -94,8 +96,7 @@ class CmdDataRemove(CmdBase): def _remove_cloud_cache(self, data_item): if not self.parsed_args.keep_in_cloud: - aws_key = self.cache_file_key(data_item.cache.dvc) - self.remove_from_cloud(aws_key) + self.cloud.remove_from_cloud(data_item) def _remove_state_file(self, data_item): if os.path.isfile(data_item.state.relative): @@ -119,24 +120,6 @@ class CmdDataRemove(CmdBase): self.remove_dir_if_empty(dvc_path.relative) Logger.debug(u'[Cmd-Remove] Remove {}. Success.'.format(name)) - def remove_from_cloud(self, aws_file_name): - Logger.debug(u'[Cmd-Remove] Remove from cloud {}.'.format(aws_file_name)) - - if not self.config.aws_access_key_id or not self.config.aws_secret_access_key: - Logger.debug('[Cmd-Remove] Unable to check cache file in the cloud') - return - conn = S3Connection(self.config.aws_access_key_id, self.config.aws_secret_access_key) - bucket_name = self.config.storage_bucket - bucket = conn.lookup(bucket_name) - if bucket: - key = bucket.get_key(aws_file_name) - if not key: - Logger.warn('[Cmd-Remove] S3 remove warning: file "{}" does not exist in S3'.format(aws_file_name)) - else: - key.delete() - Logger.info('[Cmd-Remove] File "{}" was removed from S3'.format(aws_file_name)) - pass - def remove_dir_file_by_file(self, target): for f in os.listdir(target): file = os.path.join(target, f) diff --git a/dvc/config.py b/dvc/config.py index 333651b77..aab9ae924 100644 --- a/dvc/config.py +++ b/dvc/config.py @@ -31,20 +31,6 @@ class ConfigI(object): def state_dir(self): return self._state_dir - @property - def storage_prefix(self): - return '' - - def sanity_check(self): - pass - - def aws_access_key_id(self): - return None - - @property - def aws_secret_access_key(self): - return None - class Config(ConfigI): CONFIG = 'dvc.conf' @@ -67,8 +53,6 @@ class Config(ConfigI): level = self._config['Global']['LogLevel'] Logger.set_level(level) - self._aws_creds = None - super(Config, self).__init__(self._config['Global']['DataDir'], self._config['Global']['CacheDir'], self._config['Global']['StateDir']) @@ -77,173 +61,3 @@ class Config(ConfigI): @property def file(self): return self._conf_file - - @property - def aws_access_key_id(self): - if not self._aws_creds: - self._aws_creds = self.get_aws_credentials() - if not self._aws_creds: - return None - return self._aws_creds[0] - - @property - def aws_secret_access_key(self): - if not self._aws_creds: - self._aws_creds = self.get_aws_credentials() - if not self._aws_creds: - return None - return self._aws_creds[1] - - @property - def aws_region_host(self): - """ get the region host needed for s3 access - - See notes http://docs.aws.amazon.com/general/latest/gr/rande.html#s3_region - """ - - region = self._config['AWS']['Region'] - if region is None or region == '': - return 's3.amazonaws.com' - if region == 'us-east-1': - return 's3.amazonaws.com' - return 's3.%s.amazonaws.com' % region - - - @property - def storage_path(self): - """ get storage path - - Precedence: Storage, then cloud specific - """ - - path = self._config['Global'].get('StoragePath', None) - if path: - return path - - cloud = self.cloud - assert cloud in ['aws', 'gcp'], 'unknown cloud %s' % cloud - if cloud == 'aws': - path = self._config['AWS'].get('StoragePath', None) - elif cloud == 'gcp': - path = self._config['GCP'].get('StoragePath', None) - if path is None: - raise ConfigError('invalid StoragePath: not set for Data or cloud specific') - return path - - def _storage_path_parts(self): - return self.storage_path.strip('/').split('/', 1) - - @property - def storage_bucket(self): - """ Data -> StoragePath takes precedence; if doesn't exist, use cloud-specific """ - return self._storage_path_parts()[0] - - @property - def storage_prefix(self): - parts = self._storage_path_parts() - if len(parts) > 1: - return parts[1] - return '' - - @property - def gc_project_name(self): - return self._config['GCP']['ProjectName'] - - @property - def cloud(self): - """ get cloud choice: currently one of ['AWS', 'GCP'] """ - conf = self._config['Global'].get('Cloud', '') - if conf.lower() in ['aws']: - return 'aws' - if conf.lower() in ['gcp']: - return 'gcp' - return conf - - - def get_aws_credentials(self): - """ gets aws credentials, looking in various places - - Params: - - Searches: - 1 any override in dvc.conf [AWS] CredentialPath; - 2 ~/.aws/credentials - - - Returns: - if successfully found, (access_key_id, secret) - None otherwise - """ - default = os.path.expanduser('~/.aws/credentials') - - paths = [] - credpath = self._config['AWS'].get('CredentialPath', None) - if credpath is not None and len(credpath) > 0: - credpath = os.path.expanduser(credpath) - if os.path.isfile(credpath): - paths.append(credpath) - else: - Logger.warn('AWS CredentialPath "%s" not found; falling back to default "%s"' % (credpath, default)) - paths.append(default) - else: - paths.append(default) - - for path in paths: - cc = configparser.SafeConfigParser() - threw = False - try: - # use readfp(open( ... to aid mocking. - cc.readfp(open(path, 'r')) - except Exception as e: - threw = True - if not threw and 'default' in cc.keys(): - access_key = cc['default'].get('aws_access_key_id', None) - secret = cc['default'].get('aws_secret_access_key', None) - - if access_key is not None and secret is not None: - return (access_key, secret) - - return None - - def sanity_check(self): - """ sanity check a config - - check that we have a cloud and storagePath - if aws, check can read credentials - if google, check ProjectName - - Returns: - (T,) if good - (F, issues) if bad - """ - errors = [] - for key in ['Cloud']: - if key.lower() not in self._config['Global'].keys() or len(self._config['Global'][key]) < 1: - errors.append('Please set %s in section Global in config file %s' % (key, self.file)) - - # now that a cloud is chosen, can check StoragePath - sp = self.storage_path - if sp is None or len(sp) == 0: - errors.append('Please set StoragePath = bucket/{optional path} in conf file "%s" ' - 'either in Global or a cloud specific section' % self.CONFIG) - - cloud = self.cloud - if cloud == '': - # not set; already complained above - pass - elif cloud == 'aws': - creds = self.get_aws_credentials() - if creds is None: - errors.append('can\'t find aws credentials.') - self._aws_creds = creds - elif cloud == 'gcp': - project = self._config['GCP'].get('ProjectName', None) - if project is None or len(project) < 1: - errors.append('can\'t read google cloud project name. Please set ProjectName in section GC.') - else: - errors.append('unknown Cloud %s' % cloud) - - return (len(errors) == 0, errors) - - - diff --git a/dvc/data_cloud.py b/dvc/data_cloud.py new file mode 100644 index 000000000..7491c7f57 --- /dev/null +++ b/dvc/data_cloud.py @@ -0,0 +1,379 @@ +import base64 +import hashlib +import os +import configparser +import shutil + +from boto.s3.connection import S3Connection +from google.cloud import storage as gc + +from dvc.command.base import CmdBase, DvcLock +from dvc.logger import Logger +from dvc.exceptions import DvcException +from dvc.runtime import Runtime +from dvc.system import System +from dvc.config import ConfigError + +class DataCloudError(DvcException): + def __init__(self, msg): + super(DataCloudError, self).__init__('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 DataCloudBase(object): + """ Base class for DataCloud """ + def __init__(self, settings, config, cloud_config): + self._settings = settings + self._config = config + self._cloud_config = cloud_config + + @property + def storage_path(self): + """ get storage path + + Precedence: Storage, then cloud specific + """ + + path = self._config['Global'].get('StoragePath', None) + if path: + return path + + path = self._cloud_config.get('StoragePath', None) + if path is None: + raise ConfigError('invalid StoragePath: not set for Data or cloud specific') + + return path + + def _storage_path_parts(self): + return self.storage_path.strip('/').split('/', 1) + + @property + def storage_bucket(self): + """ Data -> StoragePath takes precedence; if doesn't exist, use cloud-specific """ + return self._storage_path_parts()[0] + + @property + def storage_prefix(self): + parts = self._storage_path_parts() + if len(parts) > 1: + return parts[1] + return '' + + def cache_file_key(self, file): + return '{}/{}'.format(self.storage_prefix, file).strip('/') + + def sanity_check(self): + pass + + def sync_to_cloud(self, item): + pass + + def sync_from_cloud(self, item): + pass + + def sync(self, fname): + item = self._settings.path_factory.data_item(fname) + if os.path.isfile(item.resolved_cache.dvc): + self.sync_to_cloud(item) + else: + self.sync_from_cloud(item) + + def remove_from_cloud(self, item): + pass + + +class DataCloudLOCAL(DataCloudBase): + def sync_to_cloud(self, item): + Logger.debug('sync to cloud ' + item.resolved_cache.dvc + " " + self.storage_path) + shutil.copy(item.resolved_cache.dvc, self.storage_path) + + def sync_from_cloud(self, item): + Logger.debug('sync from cloud ' + self.storage_path + " " + item.resolved_cache.dvc) + shutil.copy(self.storage_path, item.resolved_cache.dvc) + + def remove_from_cloud(self, item): + Logger.debug('rm from cloud ' + item.resolved_cache.dvc) + os.remove(item.resolved_cache.dvc) + +class DataCloudAWS(DataCloudBase): + """ DataCloud class for Amazon Web Services """ + @property + def aws_access_key_id(self): + if not self._aws_creds: + self._aws_creds = self.get_aws_credentials() + if not self._aws_creds: + return None + return self._aws_creds[0] + + @property + def aws_secret_access_key(self): + if not self._aws_creds: + self._aws_creds = self.get_aws_credentials() + if not self._aws_creds: + return None + return self._aws_creds[1] + + @property + def aws_region_host(self): + """ get the region host needed for s3 access + + See notes http://docs.aws.amazon.com/general/latest/gr/rande.html#s3_region + """ + + region = self._cloud_config['Region'] + if region is None or region == '': + return 's3.amazonaws.com' + if region == 'us-east-1': + return 's3.amazonaws.com' + return 's3.%s.amazonaws.com' % region + + def get_aws_credentials(self): + """ gets aws credentials, looking in various places + + Params: + + Searches: + 1 any override in dvc.conf [AWS] CredentialPath; + 2 ~/.aws/credentials + + + Returns: + if successfully found, (access_key_id, secret) + None otherwise + """ + default = os.path.expanduser('~/.aws/credentials') + + paths = [] + credpath = self._cloud_config.get('CredentialPath', None) + if credpath is not None and len(credpath) > 0: + credpath = os.path.expanduser(credpath) + if os.path.isfile(credpath): + paths.append(credpath) + else: + Logger.warn('AWS CredentialPath "%s" not found; falling back to default "%s"' % (credpath, default)) + paths.append(default) + else: + paths.append(default) + + for path in paths: + cc = configparser.SafeConfigParser() + threw = False + try: + # use readfp(open( ... to aid mocking. + cc.readfp(open(path, 'r')) + except Exception as e: + threw = True + if not threw and 'default' in cc.keys(): + access_key = cc['default'].get('aws_access_key_id', None) + secret = cc['default'].get('aws_secret_access_key', None) + + if access_key is not None and secret is not None: + return (access_key, secret) + + return None + + def sanity_check(self): + creds = self.get_aws_credentials() + if creds is None: + raise ConfigError('can\'t find aws credentials.') + self._aws_creds = creds + + def _get_bucket_aws(self): + """ get a bucket object, aws """ + conn = S3Connection(self.aws_access_key_id, self.aws_secret_access_key, host=self.aws_region_host) + bucket_name = self.storage_bucket + bucket = conn.lookup(bucket_name) + if bucket is None: + raise DataCloudError('Storage path is not setup correctly') + return bucket + + def sync_from_cloud(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 DataCloudError('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(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 remove_from_cloud(self, data_item): + aws_file_name = self.cache_file_key(data_item.cache.dvc) + + Logger.debug(u'[Cmd-Remove] Remove from cloud {}.'.format(aws_file_name)) + + if not self.aws_access_key_id or not self.aws_secret_access_key: + Logger.debug('[Cmd-Remove] Unable to check cache file in the cloud') + return + conn = S3Connection(self.aws_access_key_id, self.aws_secret_access_key) + bucket_name = self.storage_bucket + bucket = conn.lookup(bucket_name) + if bucket: + key = bucket.get_key(aws_file_name) + if not key: + Logger.warn('[Cmd-Remove] S3 remove warning: file "{}" does not exist in S3'.format(aws_file_name)) + else: + key.delete() + Logger.info('[Cmd-Remove] File "{}" was removed from S3'.format(aws_file_name)) + pass + + +class DataCloudGCP(DataCloudBase): + """ DataCloud class for Google Cloud Platform """ + @property + def gc_project_name(self): + return self._cloud_config.get('ProjectName', None) + + def sanity_check(self): + project = self.gc_project_name + if project is None or len(project) < 1: + raise ConfigError('can\'t read google cloud project name. Please set ProjectName in section GC.') + + def _get_bucket_gc(self): + """ get a bucket object, gc """ + client = gc.Client(project=self.gc_project_name) + bucket = client.bucket(self.storage_bucket) + if not bucket.exists(): + raise DataCloudError('sync up: google cloud bucket {} doesn\'t exist'.format(self.storage_bucket)) + return bucket + + def sync_from_cloud(self, item): + """ sync from cloud, gcp version """ + + bucket = self._get_bucket_gc() + key = self.cache_file_key(item.resolved_cache.dvc) + + blob = bucket.get_blob(key) + if not blob: + raise DataCloudError('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(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 remove_from_cloud(self, item): + raise Exception('NOT IMPLEMENTED YET') + + +class DataCloud(object): + """ Generic class to do initial config parsing and redirect to proper DataCloud methods """ + + CLOUD_MAP = { + 'AWS': DataCloudAWS, + 'GCP': DataCloudGCP, + 'LOCAL': DataCloudLOCAL, + } + + def __init__(self, settings): + #To handle ConfigI case + if not hasattr(settings.config, '_config'): + self._cloud = DataCloudBase(None, None, None) + return + + self._settings = settings + self._config = self._settings.config._config + + self.typ = self._config['Global'].get('Cloud', '').strip().upper() + if self.typ not in self.CLOUD_MAP.keys(): + raise ConfigError('Wrong cloud type %s specified' % self.typ) + + self._cloud = self.CLOUD_MAP[self.typ](self._settings, self._config, self._config[self.typ]) + + self.sanity_check() + + def sanity_check(self): + """ sanity check a config + + check that we have a cloud and storagePath + if aws, check can read credentials + if google, check ProjectName + + Returns: + (T,) if good + (F, issues) if bad + """ + for key in ['Cloud']: + if key.lower() not in self._config['Global'].keys() or len(self._config['Global'][key]) < 1: + raise ConfigError('Please set %s in section Global in config file %s' % (key, self.file)) + + # now that a cloud is chosen, can check StoragePath + sp = self._cloud.storage_path + if sp is None or len(sp) == 0: + raise ConfigError('Please set StoragePath = bucket/{optional path} in conf file "%s" ' + 'either in Global or a cloud specific section' % self.CONFIG) + + self._cloud.sanity_check() + + def sync(self, fname): + return self._cloud.sync(fname) + + def remove_from_cloud(self, item): + return self._cloud.remove_from_cloud(item) diff --git a/tests/test_config.py b/tests/test_config.py index 73d9dc51d..f586c5260 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -10,7 +10,7 @@ import mock from dvc.path.path import Path from dvc.system import System from dvc.config import Config - +from dvc.command.data_cloud import DataCloud class TestConfigTest(TestCase): def setUp(self): @@ -24,20 +24,20 @@ class TestConfigTest(TestCase): "DataDir = ", "CacheDir = ", "StateDir = ", - "Cloud = amazon", + "Cloud = aws", "StoragePath = globalsb/global_storage_path", "[AWS]", "StoragePath = awssb/aws_storage_path", "CredentialPath =", - "[GC]", + "[GCP]", "StoragePath = googlesb/google_storage_path" ) s = StringIO.StringIO('\n'.join(c)) conf = Config(s, conf_pseudo_file = s) - - self.assertEqual(conf.cloud, 'amazon') - self.assertEqual(conf.storage_bucket, 'globalsb') - self.assertEqual(conf.storage_prefix, 'global_storage_path') + cloud = DataCloud(conf) + self.assertEqual(cloud.typ, 'aws') + self.assertEqual(cloud._cloud.storage_bucket, 'globalsb') + self.assertEqual(cloud._cloud.storage_prefix, 'global_storage_path') c = ("[Global]", "LogLevel =", @@ -48,15 +48,15 @@ class TestConfigTest(TestCase): "[AWS]", "StoragePath = awssb/aws_storage_path", "CredentialPath =", - "[GC]", + "[GCP]", "StoragePath = googlesb/google_storage_path" ) s = StringIO.StringIO('\n'.join(c)) conf = Config(s, conf_pseudo_file=s) - - self.assertEqual(conf.cloud, 'aws') - self.assertEqual(conf.storage_bucket, 'awssb') - self.assertEqual(conf.storage_prefix, 'aws_storage_path') + cloud = DataCloud(conf) + self.assertEqual(cloud.typ, 'aws') + self.assertEqual(cloud._cloud.storage_bucket, 'awssb') + self.assertEqual(cloud._cloud.storage_prefix, 'aws_storage_path') def mocked_open_aws_default_credentials(klass, file, asdf): @@ -82,7 +82,7 @@ class TestConfigTest(TestCase): "DataDir = ", "CacheDir = ", "StateDir = ", - "Cloud = amazon", + "Cloud = aws", "[AWS]", "StoragePath = awssb/aws_storage_path", "CredentialPath =", @@ -92,7 +92,8 @@ class TestConfigTest(TestCase): patcher = mock.patch('__builtin__.open', side_effect=self.mocked_open_aws_default_credentials) patcher.start() conf = Config(s, conf_pseudo_file=s) - aws_creds = conf.get_aws_credentials() + cloud = DataCloud(conf) + aws_creds = cloud._cloud.get_aws_credentials() patcher.stop() self.assertEqual(aws_creds[0], 'default_access_id') @@ -112,7 +113,7 @@ class TestConfigTest(TestCase): "DataDir = ", "CacheDir = ", "StateDir = ", - "Cloud = amazon", + "Cloud = aws", "[AWS]", "StoragePath = awssb/aws_storage_path", "CredentialPath = some_credential_path", @@ -122,7 +123,8 @@ class TestConfigTest(TestCase): patcher = mock.patch('__builtin__.open', side_effect=self.mocked_open_aws_default_credentials) patcher.start() conf = Config(s, conf_pseudo_file=s) - aws_creds = conf.get_aws_credentials() + cloud = DataCloud(conf) + aws_creds = cloud._cloud.get_aws_credentials() patcher.stop() self.assertEqual(aws_creds[0], 'override_access_id')