import: add resume flag (#1540)

Resume import

dvc/command/imp.py

@@ -15,7 +15,7 @@ class CmdImport(CmdBase):
 
             out = self.args.out or default_out
 
-            self.project.imp(self.args.url, out)
+            self.project.imp(self.args.url, out, self.args.resume)
         except DvcException:
             logger.error(
                 "failed to import {}. You could also try downloading "
@@ -48,5 +48,11 @@ def add_parser(subparsers, parent_parser):
         "'remote://myremote/path/to/file'(see "
         "`dvc remote` commands). ",
     )
+    import_parser.add_argument(
+        "--resume",
+        action="store_true",
+        default=False,
+        help="Resume previously started download.",
+    )
     import_parser.add_argument("out", nargs="?", help="Output.")
     import_parser.set_defaults(func=CmdImport)

dvc/output/base.py

@@ -164,8 +164,8 @@ class OutputBase(object):
 
         return ret
 
-    def download(self, to_info):
-        self.remote.download([self.path_info], [to_info])
+    def download(self, to_info, resume=False):
+        self.remote.download([self.path_info], [to_info], resume=resume)
 
     def checkout(self, force=False):
         if not self.use_cache:

dvc/project.py

@@ -413,7 +413,7 @@ class Project(object):
 
         return stage
 
-    def imp(self, url, out):
+    def imp(self, url, out, resume=False):
         from dvc.stage import Stage
 
         stage = Stage.create(project=self, cmd=None, deps=[url], outs=[out])
@@ -425,7 +425,7 @@ class Project(object):
 
         self.files_to_git_add = []
         with self.state:
-            stage.run()
+            stage.run(resume=resume)
 
         stage.dump()
 

dvc/remote/azure.py

@@ -145,7 +145,12 @@ class RemoteAzure(RemoteBase):
                 progress.finish_target(name)
 
     def download(
-        self, from_infos, to_infos, no_progress_bar=False, names=None
+        self,
+        from_infos,
+        to_infos,
+        no_progress_bar=False,
+        names=None,
+        resume=False,
     ):
         names = self._verify_path_args(from_infos, to_infos, names)
 

dvc/remote/base.py

@@ -167,7 +167,14 @@ class RemoteBase(object):
     def save(self, path_info):
         raise NotImplementedError
 
-    def download(self, from_infos, to_infos, no_progress_bar=False, name=None):
+    def download(
+        self,
+        from_infos,
+        to_infos,
+        no_progress_bar=False,
+        name=None,
+        resume=False,
+    ):
         raise NotImplementedError
 
     def upload(self, from_infos, to_infos, path_info, names=None):

dvc/remote/gs.py

@@ -162,7 +162,12 @@ class RemoteGS(RemoteBase):
             progress.finish_target(name)
 
     def download(
-        self, from_infos, to_infos, no_progress_bar=False, names=None
+        self,
+        from_infos,
+        to_infos,
+        no_progress_bar=False,
+        names=None,
+        resume=False,
     ):
         names = self._verify_path_args(from_infos, to_infos, names)
 

dvc/remote/hdfs.py

@@ -145,7 +145,12 @@ class RemoteHDFS(RemoteBase):
             self.hadoop_fs(cmd, user=to_info["user"])
 
     def download(
-        self, from_infos, to_infos, no_progress_bar=False, names=None
+        self,
+        from_infos,
+        to_infos,
+        no_progress_bar=False,
+        names=None,
+        resume=False,
     ):
         names = self._verify_path_args(from_infos, to_infos, names)
 

dvc/remote/http.py

@@ -28,7 +28,7 @@ class RemoteHTTP(RemoteBase):
     scheme = "http"
     REGEX = r"^https?://.*$"
     REQUEST_TIMEOUT = 10
-    CHUNK_SIZE = 1000000  # Megabyte
+    CHUNK_SIZE = 2 ** 16
     PARAM_CHECKSUM = "etag"
 
     def __init__(self, project, config):
@@ -42,7 +42,12 @@ class RemoteHTTP(RemoteBase):
         return self.cache_dir
 
     def download(
-        self, from_infos, to_infos, no_progress_bar=False, names=None
+        self,
+        from_infos,
+        to_infos,
+        no_progress_bar=False,
+        names=None,
+        resume=False,
     ):
         names = self._verify_path_args(to_infos, from_infos, names)
 
@@ -58,7 +63,6 @@ class RemoteHTTP(RemoteBase):
             )
             logger.debug(msg)
 
-            tmp_file = self.tmp_file(to_info["path"])
             if not name:
                 name = os.path.basename(to_info["path"])
 
@@ -72,14 +76,18 @@ class RemoteHTTP(RemoteBase):
                 cb = ProgressBarCallback(name, total)
 
             try:
-                self._download_to(from_info["path"], tmp_file, callback=cb)
+                self._download_to(
+                    from_info["path"],
+                    to_info["path"],
+                    callback=cb,
+                    resume=resume,
+                )
+
             except Exception:
                 msg = "failed to download '{}'".format(from_info["path"])
                 logger.error(msg)
                 continue
 
-            os.rename(tmp_file, to_info["path"])
-
             if not no_progress_bar:
                 progress.finish_target(name)
 
@@ -125,18 +133,54 @@ class RemoteHTTP(RemoteBase):
 
         return etag
 
-    def _download_to(self, url, file, callback=None):
-        r = self._request("GET", url, stream=True)
+    def _download_to(self, url, target_file, callback=None, resume=False):
+        request = self._request("GET", url, stream=True)
+        partial_file = target_file + ".part"
 
-        with open(file, "wb") as fd:
-            bytes_transfered = 0
+        mode, transferred_bytes = self._determine_mode_get_transferred_bytes(
+            partial_file, resume
+        )
 
-            for chunk in r.iter_content(chunk_size=self.CHUNK_SIZE):
-                fd.write(chunk)
+        self._validate_existing_file_size(transferred_bytes, partial_file)
+
+        self._write_request_content(
+            mode, partial_file, request, transferred_bytes, callback
+        )
+
+        os.rename(partial_file, target_file)
+
+    def _write_request_content(
+        self, mode, partial_file, request, transferred_bytes, callback=None
+    ):
+        with open(partial_file, mode) as fd:
+
+            for index, chunk in enumerate(
+                request.iter_content(chunk_size=self.CHUNK_SIZE)
+            ):
+                chunk_number = index + 1
+                if chunk_number * self.CHUNK_SIZE > transferred_bytes:
+                    fd.write(chunk)
+                    fd.flush()
+                    transferred_bytes += len(chunk)
 
                 if callback:
-                    bytes_transfered += len(chunk)
-                    callback(bytes_transfered)
+                    callback(transferred_bytes)
+
+    def _validate_existing_file_size(self, bytes_transferred, partial_file):
+        if bytes_transferred % self.CHUNK_SIZE != 0:
+            raise DvcException(
+                "File {}, might be corrupted, please remove "
+                "it and retry importing".format(partial_file)
+            )
+
+    def _determine_mode_get_transferred_bytes(self, partial_file, resume):
+        if os.path.exists(partial_file) and resume:
+            mode = "ab"
+            bytes_transfered = os.path.getsize(partial_file)
+        else:
+            mode = "wb"
+            bytes_transfered = 0
+        return mode, bytes_transfered
 
     def _request(self, method, url, **kwargs):
         kwargs.setdefault("allow_redirects", True)

dvc/remote/local.py

@@ -487,7 +487,12 @@ class RemoteLOCAL(RemoteBase):
                 )
 
     def download(
-        self, from_infos, to_infos, no_progress_bar=False, names=None
+        self,
+        from_infos,
+        to_infos,
+        no_progress_bar=False,
+        names=None,
+        resume=False,
     ):
         names = self._verify_path_args(from_infos, to_infos, names)
 

dvc/remote/s3.py

@@ -207,7 +207,12 @@ class RemoteS3(RemoteBase):
             progress.finish_target(name)
 
     def download(
-        self, from_infos, to_infos, no_progress_bar=False, names=None
+        self,
+        from_infos,
+        to_infos,
+        no_progress_bar=False,
+        names=None,
+        resume=False,
     ):
         names = self._verify_path_args(from_infos, to_infos, names)
 

dvc/remote/ssh.py

@@ -279,7 +279,12 @@ class RemoteSSH(RemoteBase):
         ssh.close()
 
     def download(
-        self, from_infos, to_infos, no_progress_bar=False, names=None
+        self,
+        from_infos,
+        to_infos,
+        no_progress_bar=False,
+        names=None,
+        resume=False,
     ):
         names = self._verify_path_args(from_infos, to_infos, names)
 

dvc/stage.py

@@ -564,7 +564,7 @@ class Stage(object):
         if p.returncode != 0:
             raise StageCmdFailedError(self)
 
-    def run(self, dry=False):
+    def run(self, dry=False, resume=False):
         if self.locked:
             logger.info(
                 "Verifying outputs in locked stage '{stage}'".format(
@@ -584,7 +584,9 @@ class Stage(object):
                 if self._already_cached():
                     self.outs[0].checkout()
                 else:
-                    self.deps[0].download(self.outs[0].path_info)
+                    self.deps[0].download(
+                        self.outs[0].path_info, resume=resume
+                    )
 
         elif self.is_data_source:
             msg = "Verifying data sources in '{}'".format(self.relpath)

tests/test_import.py

@@ -4,10 +4,12 @@ import os
 from uuid import uuid4
 
 from dvc import logger
+from dvc.utils.compat import urljoin
 from dvc.exceptions import DvcException
 from dvc.main import main
-from mock import patch
+from mock import patch, mock_open, call
 from tests.basic_env import TestDvc
+from tests.utils.httpd import StaticFileServer
 
 from dvc.utils.compat import StringIO
 
@@ -66,3 +68,72 @@ class TestFailedImportMessage(TestDvc):
             "command.",
             logger.logger.handlers[1].stream.getvalue(),
         )
+
+
+class TestInterruptedDownload(TestDvc):
+    @property
+    def remote(self):
+        return "http://localhost:8000/"
+
+    def _prepare_interrupted_download(self):
+        import_url = urljoin(self.remote, self.FOO)
+        import_output = "imported_file"
+        tmp_file_name = import_output + ".part"
+        tmp_file_path = os.path.join(self._root_dir, tmp_file_name)
+        self._import_with_interrupt(import_output, import_url)
+        self.assertTrue(os.path.exists(tmp_file_name))
+        self.assertFalse(os.path.exists(import_output))
+        return import_output, import_url, tmp_file_path
+
+    def _import_with_interrupt(self, import_output, import_url):
+        def interrupting_generator():
+            yield self.FOO[0].encode("utf8")
+            raise KeyboardInterrupt
+
+        with patch(
+            "requests.models.Response.iter_content",
+            return_value=interrupting_generator(),
+        ):
+            with patch(
+                "dvc.remote.http.RemoteHTTP._content_length", return_value=3
+            ):
+                result = main(["import", import_url, import_output])
+                self.assertEqual(result, 252)
+
+
+class TestShouldResumeDownload(TestInterruptedDownload):
+    @patch("dvc.remote.http.RemoteHTTP.CHUNK_SIZE", 1)
+    def test(self):
+        with StaticFileServer():
+            import_output, import_url, tmp_file_path = (
+                self._prepare_interrupted_download()
+            )
+
+            m = mock_open()
+            with patch("dvc.remote.http.open", m):
+                result = main(
+                    ["import", "--resume", import_url, import_output]
+                )
+                self.assertEqual(result, 0)
+        m.assert_called_once_with(tmp_file_path, "ab")
+        m_handle = m()
+        expected_calls = [call(b"o"), call(b"o")]
+        m_handle.write.assert_has_calls(expected_calls, any_order=False)
+
+
+class TestShouldNotResumeDownload(TestInterruptedDownload):
+    @patch("dvc.remote.http.RemoteHTTP.CHUNK_SIZE", 1)
+    def test(self):
+        with StaticFileServer():
+            import_output, import_url, tmp_file_path = (
+                self._prepare_interrupted_download()
+            )
+
+            m = mock_open()
+            with patch("dvc.remote.http.open", m):
+                result = main(["import", import_url, import_output])
+                self.assertEqual(result, 0)
+        m.assert_called_once_with(tmp_file_path, "wb")
+        m_handle = m()
+        expected_calls = [call(b"f"), call(b"o"), call(b"o")]
+        m_handle.write.assert_has_calls(expected_calls, any_order=False)

tests/utils/httpd.py

@@ -1,5 +1,6 @@
 import hashlib
 import os
+import socket
 import threading
 
 from dvc.utils.compat import HTTPServer, SimpleHTTPRequestHandler
@@ -30,13 +31,26 @@ class ContentMD5Handler(SimpleHTTPRequestHandler):
 
 
 class StaticFileServer:
-    __server_lock = threading.Lock()
+    _server_lock = threading.Lock()
 
     def __init__(self, handler="etag"):
+        self._server_lock.acquire()
         handler_class = ETagHandler if handler == "etag" else ContentMD5Handler
+        self.bind_port_by_any_means(handler_class)
 
-        self.__server_lock.acquire()
-        self.httpd = HTTPServer(("localhost", 8000), handler_class)
+    def bind_port_by_any_means(self, handler_class):
+        import time
+
+        # shutdowning/closing socket/server does not unbind port in time,
+        # locking the server also does not bring results, hence
+        # this method
+        for i in range(10000):
+            try:
+                self.httpd = HTTPServer(("localhost", 8000), handler_class)
+            except Exception:
+                time.sleep(0.01)
+                continue
+            break
 
     def __enter__(self):
         self.server_thread = threading.Thread(target=self.httpd.serve_forever)
@@ -44,6 +58,7 @@ class StaticFileServer:
         self.server_thread.start()
 
     def __exit__(self, *args):
+        self.httpd.socket.close()
         self.httpd.shutdown()
         self.httpd.server_close()
-        self.__server_lock.release()
+        self._server_lock.release()