From fc9e0912864e19c15a6886f8f9ec42b4b39c47cf Mon Sep 17 00:00:00 2001 From: Ophir Yokatan Date: Fri, 14 Jul 2017 00:17:43 +0300 Subject: [PATCH] download file to temp file, and renmae only if success --- dvc/data_cloud.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/dvc/data_cloud.py b/dvc/data_cloud.py index c99e194d0..1e245f3f7 100644 --- a/dvc/data_cloud.py +++ b/dvc/data_cloud.py @@ -3,6 +3,7 @@ import hashlib import os import threading import configparser +import tempfile from boto.s3.connection import S3Connection from google.cloud import storage as gc @@ -251,10 +252,14 @@ class DataCloudAWS(DataCloudBase): Logger.info('Downloading cache file from S3 "{}/{}"'.format(bucket.name, key_name)) try: - key.get_contents_to_filename(item.resolved_cache.relative, + temp_file = tempfile.NamedTemporaryFile(dir=item.resolved_cache.dirname, delete=False) + key.get_contents_to_filename(temp_file.name, cb=create_cb(item.resolved_cache.relative)) + os.rename(temp_file.name, item.resolved_cache.relative) except Exception as exc: Logger.error('Failed to download "{}": {}'.format(key_name, exc)) + if temp_file and os.path.exists(temp_file.name): + os.remove(temp_file.name) return progress.finish_target(os.path.basename(item.resolved_cache.relative))