From ebeb32a8acbcf05b33efdcc165872a0f3b7b4246 Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <kupruser@gmail.com>
Date: Mon, 14 Aug 2017 00:53:24 +0300
Subject: [PATCH] import: resume download

Download to .part file first and rename at the very end,
so if we lose connection we could resume download.
Note: not all servers support 'Range' header, so this
feature is not going to work everywhere.

This feature is activated by -c/--continue cmdline option.

Fixes #108

Signed-off-by: Ruslan Kuprieiev <kupruser@gmail.com>
---
 dvc/cli.py                 |  6 ++++++
 dvc/command/import_file.py | 34 ++++++++++++++++++++++++++++++----
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/dvc/cli.py b/dvc/cli.py
index 742b11aa9..10701d03f 100644
--- a/dvc/cli.py
+++ b/dvc/cli.py
@@ -223,6 +223,12 @@ def parse_args(argv=None):
                         type=int,
                         default=cpu_count(),
                         help='Number of jobs to run simultaneously.')
+    import_parser.add_argument('-c',
+                        '--continue',
+                        dest='cont',
+                        action='store_true',
+                        default=False,
+                        help='Resume downloading file from url')
     import_parser.set_defaults(func=CmdImportFile)
 
     # Lock
diff --git a/dvc/command/import_file.py b/dvc/command/import_file.py
index 201f93cc7..981fd6f4c 100644
--- a/dvc/command/import_file.py
+++ b/dvc/command/import_file.py
@@ -142,21 +142,45 @@ class CmdImportFile(CmdBase):
         """
         Download single file from url.
         """
-        r = requests.get(from_url, stream=True)
 
+        tmp_file = to_file + '.part'
         name = os.path.basename(from_url)
         chunk_size = 1024 * 100
         downloaded = 0
         last_reported = 0
         report_bucket = 100*1024*10
-        total_length = r.headers.get('content-length')
 
-        with open(to_file, 'wb') as f:
+        resume_header = None
+        mode = 'wb'
+
+        # Resume download if we can
+        if os.path.exists(tmp_file) and self.parsed_args.cont:
+            mode = 'ab'
+
+            downloaded = os.path.getsize(tmp_file)
+            resume_header = {'Range': 'bytes=%d-' % downloaded}
+
+            Logger.debug('found existing {} file, resuming download'.format(tmp_file))
+
+        r = requests.get(from_url, stream=True, headers=resume_header)
+
+        content_range = r.headers.get('content-range')
+        if resume_header and content_range == None:
+            mode = 'wb'
+            downloaded = 0
+
+            Logger.debug('\'range\' is not supported by the server. Can\'t resume download')
+
+        total_length = r.headers.get('content-length')
+        if total_length == None:
+            Logger.debug('\'content-length\' is not supported by the server')
+
+        with open(tmp_file, mode) as f:
             for chunk in r.iter_content(chunk_size=chunk_size):
                 if not chunk:  # filter out keep-alive new chunks
                     continue
 
-                downloaded += chunk_size
+                downloaded += len(chunk)
 
                 last_reported += chunk_size
                 if last_reported >= report_bucket:
@@ -171,6 +195,8 @@ class CmdImportFile(CmdBase):
         # tell progress bar that this target is finished downloading
         progress.finish_target(name)
 
+        os.rename(tmp_file, to_file)
+
     def download_targets(self, targets):
         """
         Download targets in a number of threads.