diff --git a/dvc/cloud/data_cloud.py b/dvc/cloud/data_cloud.py index baee0b3a7..050cbfd12 100644 --- a/dvc/cloud/data_cloud.py +++ b/dvc/cloud/data_cloud.py @@ -107,13 +107,12 @@ class DataCloud(object): def _filter(self, func, status, targets, jobs, remote): cloud = self._get_cloud(remote) - with cloud: - filtered = [] - for t, s in self._status(cloud, targets, jobs): - if s == STATUS_MODIFIED or s == status: - filtered.append(t) + filtered = [] + for t, s in self._status(cloud, targets, jobs): + if s == STATUS_MODIFIED or s == status: + filtered.append(t) - return map_progress(getattr(cloud, func), filtered, jobs) + return map_progress(getattr(cloud, func), filtered, jobs) def push(self, targets, jobs=1, remote=None): """ @@ -143,5 +142,4 @@ class DataCloud(object): """ cloud = self._get_cloud(remote) - with cloud: - return self._status(cloud, targets, jobs) + return self._status(cloud, targets, jobs) diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 2bc64d1b8..151e9f9db 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -46,6 +46,12 @@ class RemoteBase(object): def match(cls, url): return re.match(cls.REGEX, url) + @staticmethod + def tmp_file(fname): + """ Temporary name for a partial download """ + #FIXME probably better use uuid() + return fname + '.part' + def save_info(self, path_info): raise NotImplementedError @@ -55,13 +61,17 @@ class RemoteBase(object): def checkout(self, path_info, checksum_info): raise NotImplementedError - def download(self, path_info, path): + def _get_path_info(self, path): raise NotImplementedError - def upload(self, path, path_info): + def _new_path_info(self, path): raise NotImplementedError - # Old code starting from here + def download(self, path_info, path, no_progress_bar=False, name=None): + raise NotImplementedError + + def upload(self, path, path_info, name=None): + raise NotImplementedError def cache_file_key(self, fname): """ Key of a file within the bucket """ @@ -69,18 +79,6 @@ class RemoteBase(object): relpath = relpath.replace('\\', '/') return posixpath.join(self.prefix, relpath).strip('/') - def cache_key_name(self, path): - relpath = os.path.relpath(path, self.project.cache.local.cache_dir) - return relpath.replace('\\', '').replace('/', '') - - @staticmethod - def tmp_file(fname): - """ Temporary name for a partial download """ - return fname + '.part' - - def _push_key(self, key, path): - pass - def collect(self, arg): from dvc.remote.local import RemoteLOCAL @@ -95,12 +93,12 @@ class RemoteBase(object): return ret dir_path = path else: - key = self._get_key(path) - if not key: + path_info = self._get_path_info(path) + if not path_info: Logger.debug("File '{}' does not exist in the cloud".format(path)) return ret tmp = os.path.join(tempfile.mkdtemp(), os.path.basename(path)) - self._pull_key(key, tmp, no_progress_bar=True) + self.download(path_info, tmp, no_progress_bar=True) dir_path = tmp for relpath, md5 in RemoteLOCAL.get_dir_cache(dir_path).items(): @@ -109,24 +107,19 @@ class RemoteBase(object): return ret - def _cmp_checksum(self, blob, fname): - md5 = self.project.cache.local.path_to_md5(fname) - if self.project.cache.local.state.changed(fname, md5=md5): - return False - - return True - def push(self, path): - key = self._get_key(path) - if key: - Logger.debug("File '{}' already uploaded to the cloud. Validating checksum...".format(path)) - if self._cmp_checksum(key, path): - Logger.debug('File checksum matches. No uploading is needed.') - return [] - Logger.debug('Checksum mismatch. Reuploading is required.') + path_info = self._get_path_info(path) + if path_info: + Logger.debug("File '{}' already uploaded to the cloud.".format(path)) + return None - key = self._new_key(path) - return self._push_key(key, path) + path_info = self._new_path_info(path) + md5 = self.project.cache.local.path_to_md5(path) + + if self.project.cache.local.changed(md5): + return None + + self.upload(path, path_info, name=md5) def _makedirs(self, fname): dname = os.path.dirname(fname) @@ -136,51 +129,32 @@ class RemoteBase(object): if e.errno != os.errno.EEXIST: raise - def _pull_key(self, key, path, no_progress_bar=False): - """ Cloud-specific method of pulling keys """ - pass - - def _get_key(self, path): - """ Cloud-specific method of getting keys """ - pass - def pull(self, path): """ Generic method for pulling data from the cloud """ - key = self._get_key(path) - if not key: + path_info = self._get_path_info(path) + if not path_info: Logger.error("File '{}' does not exist in the cloud".format(path)) return None - return self._pull_key(key, path) + md5 = self.project.cache.local.path_to_md5(path) - def _status(self, key, path): - remote_exists = key != None - local_exists = os.path.exists(path) + self.download(path_info, path, name=md5) - diff = None - if remote_exists and local_exists: - diff = self._cmp_checksum(key, path) - - return STATUS_MAP.get((local_exists, remote_exists, diff), STATUS_UNKNOWN) + if self.project.cache.local.changed(md5): + return None def status(self, path): """ Generic method for checking data item status. """ - key = self._get_key(path) - if not key: - return STATUS_NEW + md5 = self.project.cache.local.path_to_md5(path) + path_info = self._get_path_info(path) + remote_exists = path_info != None + local_exists = os.path.exists(path) and not self.project.cache.local.changed(md5) - return self._status(key, path) + diff = None + if remote_exists and local_exists: + md5 = self.project.cache.local.path_to_md5(path) + diff = not self.project.cache.local.changed(md5) - def connect(self): - pass - - def disconnect(self): - pass - - def __enter__(self): - self.connect() - - def __exit__(self, type, value, tb): - self.disconnect() + return STATUS_MAP.get((local_exists, remote_exists, diff), STATUS_UNKNOWN) diff --git a/dvc/remote/gs.py b/dvc/remote/gs.py index af400ccdd..4683fb648 100644 --- a/dvc/remote/gs.py +++ b/dvc/remote/gs.py @@ -89,42 +89,70 @@ class RemoteGS(RemoteBase): blob.delete() - def upload(self, path, path_info): + def _get_path_info(self, path): + key = self.cache_file_key(path) + if not self.gs.bucket(self.bucket).get_blob(key): + return None + return {'scheme': 'gs', + 'bucket': self.bucket, + 'key': key} + + def _new_path_info(self, path): + key = self.cache_file_key(path) + return {'scheme': 'gs', + 'bucket': self.bucket, + 'key': key} + + def upload(self, path, path_info, name=None): if path_info['scheme'] != 'gs': raise NotImplementedError - blob = self.gs.bucket(path_info['bucket']).blob(path_info['key']) - blob.upload_from_filename(path) + Logger.debug("Uploading '{}' to '{}/{}'".format(path, + path_info['bucket'], + path_info['key'])) - def download(self, path_info, path): + if not name: + name = os.path.basename(path) + + progress.update_target(name, 0, None) + + try: + self.gs.bucket(path_info['bucket']).blob(path_info['key']).upload_from_filename(path) + except Exception as exc: + Logger.error("Failed to upload '{}' to '{}/{}'".format(path, + path_info['bucket'], + path_info['key']), exc) + return None + + progress.finish_target(name) + + return path + + def download(self, path_info, path, no_progress_bar=False, name=None): if path_info['scheme'] != 'gs': raise NotImplementedError - blob = self.gs.bucket(path_info['bucket']).get_blob(path_info['key']) - blob.download_to_filename(path) + Logger.debug("Downloading '{}/{}' to '{}'".format(path_info['bucket'], + path_info['key'], + path)) - # old code from here - def _pull_key(self, key, path, no_progress_bar=False): - self._makedirs(path) - - name = self.cache_key_name(path) tmp_file = self.tmp_file(path) - - if self._cmp_checksum(key, path): - Logger.debug('File "{}" matches with "{}".'.format(path, key.name)) - return path - - Logger.debug('Downloading cache file from gc "{}/{}"'.format(key.bucket.name, key.name)) + if not name: + name = os.path.basename(path) if not no_progress_bar: # percent_cb is not available for download_to_filename, so # lets at least update progress at keypoints(start, finish) progress.update_target(name, 0, None) + self._makedirs(path) + try: - key.download_to_filename(tmp_file) + self.gs.bucket(path_info['bucket']).get_blob(path_info['key']).download_to_filename(tmp_file) except Exception as exc: - Logger.error('Failed to download "{}": {}'.format(key.name, exc)) + Logger.error("Failed to download '{}/{}' to '{}'".format(path_info['bucket'], + path_info['key'], + path), exc) return None os.rename(tmp_file, path) @@ -132,27 +160,4 @@ class RemoteGS(RemoteBase): if not no_progress_bar: progress.finish_target(name) - Logger.debug('Downloading completed') - - return path - - def _get_key(self, path): - key_name = self.cache_file_key(path) - return self.gs.bucket(self.bucket).get_blob(key_name) - - def _new_key(self, path): - key_name = self.cache_file_key(path) - return self.gs.bucket(self.bucket).blob(key_name) - - def _push_key(self, key, path): - """ push, gcp version """ - name = self.cache_key_name(path) - - progress.update_target(name, 0, None) - - key.upload_from_filename(path) - - progress.finish_target(name) - Logger.debug('uploading %s completed' % path) - return path diff --git a/dvc/remote/local.py b/dvc/remote/local.py index 6b14a1a27..71ace5132 100644 --- a/dvc/remote/local.py +++ b/dvc/remote/local.py @@ -256,58 +256,56 @@ class RemoteLOCAL(RemoteBase): remove(path_info['path']) - def _copy(self, from_path, to_path): - dname = os.path.dirname(to_path) - if not os.path.exists(dname): - os.makedirs(dname) - - shutil.copyfile(from_path, to_path) - - def upload(self, path, path_info): - if path_info['scheme'] != 'local': - raise NotImplementedError - - self._copy(path, path_info['path']) - - def download(self, path_info, path): - if path_info['scheme'] != 'local': - raise NotImplementedError - - self._copy(path_info['path'], path) - - # old code starting from here def cache_file_key(self, path): - return os.path.relpath(os.path.abspath(path), self.project.cache.local.cache_dir) + relpath = os.path.relpath(os.path.abspath(path), self.project.cache.local.cache_dir) + return os.path.join(self.prefix, relpath) - def _get_key(self, path): - key_name = self.cache_file_key(path) - key = LocalKey(self.prefix, key_name) - if os.path.exists(key.path) and os.path.isfile(key.path): - return key - return None + def _get_path_info(self, path): + key = self.cache_file_key(path) + if not os.path.exists(key) or not os.path.isfile(key): + return None + return {'scheme': 'local', + 'path': key} - def _new_key(self, path): - key_name = self.cache_file_key(path) - key = LocalKey(self.prefix, key_name) - self._makedirs(key.path) - return key + def _new_path_info(self, path): + key = self.cache_file_key(path) + return {'scheme': 'local', + 'path': key} + + def upload(self, path, path_info, name=None): + if path_info['scheme'] != 'local': + raise NotImplementedError + + Logger.debug("Uploading '{}' to '{}'".format(path, path_info['path'])) + + if not name: + name = os.path.basename(path) + + self._makedirs(path_info['path']) + + try: + copyfile(path, path_info['path'], name=name) + except Exception as exc: + Logger.error("Failed to upload '{}' tp '{}'".format(path, path_info['path'])) + return None - def _push_key(self, key, path): - self._makedirs(key.path) - name = self.cache_key_name(path) - copyfile(path, key.path, name=name) return path - def _pull_key(self, key, path, no_progress_bar=False): + def download(self, path_info, path, no_progress_bar=False, name=None): + if path_info['scheme'] != 'local': + raise NotImplementedError + + Logger.debug("Downloading '{}' to '{}'".format(path_info['path'], path)) + + if not name: + name = os.path.basename(path) + self._makedirs(path) - - name = self.cache_key_name(path) - tmp_file = self.tmp_file(path) try: - copyfile(key.path, tmp_file, no_progress_bar=no_progress_bar, name=name) + copyfile(path_info['path'], tmp_file, no_progress_bar=no_progress_bar, name=name) except Exception as exc: - Logger.error('Failed to copy "{}": {}'.format(key.path, exc)) + Logger.error("Failed to download '{}' to '{}'".format(path_info['path'], path), exc) return None os.rename(tmp_file, path) diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py index 722a78429..fdccb073b 100644 --- a/dvc/remote/s3.py +++ b/dvc/remote/s3.py @@ -110,55 +110,72 @@ class RemoteS3(RemoteBase): except Exception: pass - def upload(self, path, path_info): + def _get_path_info(self, path): + key = self.cache_file_key(path) + try: + self.s3.Object(self.bucket, key).get() + return {'scheme': self.scheme, + 'bucket': self.bucket, + 'key': key} + except Exception: + return None + + def _new_path_info(self, path): + key = self.cache_file_key(path) + return {'scheme': self.scheme, + 'bucket': self.bucket, + 'key': key} + + def upload(self, path, path_info, name=None): if path_info['scheme'] != 's3': raise NotImplementedError - self.s3.Object(path_info['bucket'], path_info['key']).upload_file(path) + Logger.debug("Uploading '{}' to '{}/{}'".format(path, + path_info['bucket'], + path_info['key'])) + if not name: + name = os.path.basename(path) - def download(self, path_info, path): + total = os.path.getsize(path) + cb = Callback(name, total) + + try: + self.s3.Object(path_info['bucket'], path_info['key']).upload_file(path, Callback=cb) + except Exception as exc: + Logger.error("Failed to upload '{}'".format(path), exc) + return None + + progress.finish_target(name) + + return path + + def download(self, path_info, fname, no_progress_bar=False, name=None): if path_info['scheme'] != 's3': raise NotImplementedError - self.s3.Object(path_info['bucket'], path_info['key']).download_file(path) - - # Old code starting from here - def create_cb_pull(self, name, key): - total = self.s3.Object(bucket_name=key.bucket, key=key.name).content_length - return Callback(name, total) - - def create_cb_push(self, name, fname): - total = os.path.getsize(fname) - return Callback(name, total) - - def _pull_key(self, key, fname, no_progress_bar=False): - Logger.debug("Pulling key '{}' from bucket '{}' to file '{}'".format(key.name, - key.bucket, - fname)) - self._makedirs(fname) + Logger.debug("Downloading '{}/{}' to '{}'".format(path_info['bucket'], + path_info['key'], + fname)) tmp_file = self.tmp_file(fname) - name = self.cache_key_name(fname) - - if self._cmp_checksum(key, fname): - Logger.debug('File "{}" matches with "{}".'.format(fname, key.name)) - return fname - - Logger.debug('Downloading cache file from S3 "{}/{}" to "{}"'.format(key.bucket, - key.name, - fname)) + if not name: + name = os.path.basename(fname) if no_progress_bar: cb = None else: - cb = self.create_cb_pull(name, key) + total = self.s3.Object(bucket_name=path_info['bucket'], + key=path_info['key']).content_length + cb = Callback(name, total) + self._makedirs(fname) try: - self.s3.Object(key.bucket, key.name).download_file(tmp_file, Callback=cb) + self.s3.Object(path_info['bucket'], path_info['key']).download_file(tmp_file, Callback=cb) except Exception as exc: - Logger.error('Failed to download "{}": {}'.format(key.name, exc)) + Logger.error("Failed to download '{}/{}'".format(path_info['bucket'], + path_info['key']), exc) return None os.rename(tmp_file, fname) @@ -166,32 +183,4 @@ class RemoteS3(RemoteBase): if not no_progress_bar: progress.finish_target(name) - Logger.debug('Downloading completed') - return fname - - def _get_key(self, path): - key_name = self.cache_file_key(path) - try: - self.s3.Object(self.bucket, key_name).get() - return AWSKey(self.bucket, key_name) - except Exception: - return None - - def _new_key(self, path): - key_name = self.cache_file_key(path) - return AWSKey(self.bucket, key_name) - - def _push_key(self, key, path): - """ push, aws version """ - name = self.cache_key_name(path) - cb = self.create_cb_push(name, path) - try: - self.s3.Object(key.bucket, key.name).upload_file(path, Callback=cb) - except Exception as exc: - Logger.error('Failed to upload "{}": {}'.format(path, exc)) - return None - - progress.finish_target(name) - - return path diff --git a/dvc/remote/ssh.py b/dvc/remote/ssh.py index 748c316bf..785a7f6fb 100644 --- a/dvc/remote/ssh.py +++ b/dvc/remote/ssh.py @@ -49,102 +49,131 @@ class RemoteSSH(RemoteBase): #NOTE: ~/ paths are temporarily forbidden REGEX = r'^(?P.*)@(?P.*):(?P/+.*)$' - def __init__(self, project, config): self.project = project storagepath = config.get(Config.SECTION_AWS_STORAGEPATH, '/') self.url = config.get(Config.SECTION_REMOTE_URL, storagepath) + self.host = self.group('host') + self.user = self.group('user') + self.prefix = self.group('path') def group(self, group): return self.match(self.url).group(group) - @property - def path(self): - return self.group('path') + def cache_file_key(self, path): + relpath = os.path.relpath(os.path.abspath(path), self.project.cache.local.cache_dir) + return posixpath.join(self.prefix, relpath.replace('\\', '/')) - @property - def hostname(self): - return self.group('host') + def ssh(self, host=None, user=None): + Logger.debug("Establishing ssh connection with '{}' as user '{}'".format(host, user)) - @property - def username(self): - return self.group('user') - - @property - def ssh(self): ssh = paramiko.SSHClient() ssh.load_system_host_keys() ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) - ssh.connect(self.hostname, username=self.username) + ssh.connect(host, username=user) return ssh - def disconnect(self): - self.ssh.close() - - def get_sftp(self): - #NOTE: SFTP doesn't seem to be thread safe - return self.ssh.open_sftp() - - def cache_file_key(self, path): - relpath = os.path.relpath(os.path.abspath(path), self.project.cache.local.cache_dir) - return relpath.replace('\\', '/') - - def _isfile_remote(self, path): + def _isfile_remote(self, path_info): try: - self.get_sftp().open(path, 'r') + ssh = self.ssh(path_info['host'], path_info['user']) + ssh.open_sftp().open(path_info['path'], 'r') + ssh.close() return True except IOError as exc: if exc.errno != errno.ENOENT: raise return False - def _get_key(self, path): - key_name = self.cache_file_key(path) - key = SSHKey(self.path, key_name) - if self._isfile_remote(key.path): - return key + def _get_path_info(self, path): + key = self.cache_file_key(path) + ret = {'scheme': 'ssh', + 'host': self.host, + 'user': self.user, + 'path': key} + if self._isfile_remote(ret): + return ret return None - def _makedirs_remote(self, dname): + def _do_makedirs_remote(self, sftp, dname): try: - self.get_sftp().chdir(dname) + sftp.chdir(dname) except IOError as exc: if exc.errno != errno.ENOENT: raise parent = posixpath.dirname(dname) if len(parent): - self._makedirs_remote(parent) - self.get_sftp().mkdir(dname) + self._do_makedirs_remote(sftp, parent) - def _new_key(self, path): - key_name = self.cache_file_key(path) - key = SSHKey(self.path, key_name) - self._makedirs_remote(posixpath.dirname(key.path)) - return key + sftp.mkdir(dname) - def _push_key(self, key, path): - name = self.cache_key_name(path) - self.get_sftp().put(path, key.path, callback=create_cb(name)) - progress.finish_target(key.name) - return path + def _makedirs_remote(self, path_info): + dname = posixpath.dirname(path_info['path']) + ssh = self.ssh(path_info['host'], path_info['user']) + sftp = ssh.open_sftp() + self._do_makedirs_remote(sftp, dname) + ssh.close() + + def _new_path_info(self, path): + key = self.cache_file_key(path) + ret = {'scheme': 'ssh', + 'host': self.host, + 'user': self.user, + 'path': key} + self._makedirs_remote(ret) + return ret + + def download(self, path_info, path, no_progress_bar=False, name=None): + if path_info['scheme'] != 'ssh': + raise NotImplementedError + + Logger.debug("Downloading '{}/{}' to '{}'".format(path_info['host'], + path_info['path'], + path)) + if not name: + name = os.path.basename(path) - def _pull_key(self, key, path, no_progress_bar=False): self._makedirs(path) - - name = self.cache_key_name(path) - tmp_file = self.tmp_file(path) try: - self.get_sftp().get(key.path, tmp_file, callback=create_cb(name)) + ssh = self.ssh(host=path_info['host'], user=path_info['user']) + ssh.open_sftp().get(path_info['path'], tmp_file, callback=create_cb(name)) + ssh.close() except Exception as exc: - Logger.error('Failed to copy "{}": {}'.format(key.path, exc)) + Logger.error("Failed to download '{}/{}' to '{}'".format(path_info['host'], + path_info['path'], + path), exc) return None os.rename(tmp_file, path) - progress.finish_target(key.name) + progress.finish_target(name) + + return path + + def upload(self, path, path_info, name=None): + if path_info['scheme'] != 'ssh': + raise NotImplementedError + + Logger.debug("Uploading '{}' to '{}/{}'".format(path, + path_info['host'], + path_info['path'])) + + if not name: + name = os.path.basename(path) + + try: + ssh = self.ssh(host=path_info['host'], user=path_info['user']) + ssh.open_sftp().put(path, path_info['path'], callback=create_cb(name)) + ssh.close() + except Exception as exc: + Logger.error("Failed to upload '{}' to '{}/{}'".format(path, + path_info['host'], + path_info['path']), exc) + return None + + progress.finish_target(name) return path diff --git a/tests/test_data_cloud.py b/tests/test_data_cloud.py index cfa342e33..80da8601a 100644 --- a/tests/test_data_cloud.py +++ b/tests/test_data_cloud.py @@ -143,8 +143,6 @@ class TestDataCloudBase(TestDvc): def _test_cloud(self): self._setup_cloud() - self.cloud.connect() - stage = self.dvc.add(self.FOO) cache = stage.outs[0].cache @@ -179,7 +177,10 @@ class TestDataCloudBase(TestDvc): with open(cache, 'a') as fd: fd.write('addon') status = self.cloud.status(cache) - self.assertEqual(status, STATUS_MODIFIED) + # NOTE: this is what corrupted cache looks like and it should + # be removed automatically, thus causing STATUS_DELETED instead + # of STATUS_MODIFIED. + self.assertEqual(status, STATUS_DELETED) # Remove and check status sleep()