From a92bfd272dc8bcb73ea26ab4a9cd1daba73cddf6 Mon Sep 17 00:00:00 2001
From: Batuhan Taskaya <batuhanosmantaskaya@gmail.com>
Date: Tue, 2 Feb 2021 16:49:58 +0300
Subject: [PATCH] get, get-url: add --jobs option (#5362)

---
 dvc/command/get.py                 | 12 ++++++++++++
 dvc/command/get_url.py             | 13 ++++++++++++-
 dvc/repo/get.py                    |  6 ++++--
 dvc/repo/get_url.py                |  4 ++--
 tests/unit/command/test_get.py     | 16 ++++++++++++++--
 tests/unit/command/test_get_url.py |  4 ++--
 6 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/dvc/command/get.py b/dvc/command/get.py
index 7f27e34fe..5f0666f62 100644
--- a/dvc/command/get.py
+++ b/dvc/command/get.py
@@ -39,6 +39,7 @@ class CmdGet(CmdBaseNoRepo):
                 path=self.args.path,
                 out=self.args.out,
                 rev=self.args.rev,
+                jobs=self.args.jobs,
             )
             return 0
         except DvcException:
@@ -84,4 +85,15 @@ def add_parser(subparsers, parent_parser):
         help="Print the storage location (URL) the target data would be "
         "downloaded from, and exit.",
     )
+    get_parser.add_argument(
+        "-j",
+        "--jobs",
+        type=int,
+        help=(
+            "Number of jobs to run simultaneously. "
+            "The default value is 4 * cpu_count(). "
+            "For SSH remotes, the default is 4. "
+        ),
+        metavar="<number>",
+    )
     get_parser.set_defaults(func=CmdGet)
diff --git a/dvc/command/get_url.py b/dvc/command/get_url.py
index 015e4119b..6c950673b 100644
--- a/dvc/command/get_url.py
+++ b/dvc/command/get_url.py
@@ -14,7 +14,7 @@ class CmdGetUrl(CmdBaseNoRepo):
         from dvc.repo import Repo
 
         try:
-            Repo.get_url(self.args.url, out=self.args.out)
+            Repo.get_url(self.args.url, out=self.args.out, jobs=self.args.jobs)
             return 0
         except DvcException:
             logger.exception(f"failed to get '{self.args.url}'")
@@ -36,4 +36,15 @@ def add_parser(subparsers, parent_parser):
     get_parser.add_argument(
         "out", nargs="?", help="Destination path to put data to.",
     ).complete = completion.DIR
+    get_parser.add_argument(
+        "-j",
+        "--jobs",
+        type=int,
+        help=(
+            "Number of jobs to run simultaneously. "
+            "The default value is 4 * cpu_count(). "
+            "For SSH remotes, the default is 4. "
+        ),
+        metavar="<number>",
+    )
     get_parser.set_defaults(func=CmdGetUrl)
diff --git a/dvc/repo/get.py b/dvc/repo/get.py
index d2c5efb92..6b793c619 100644
--- a/dvc/repo/get.py
+++ b/dvc/repo/get.py
@@ -17,7 +17,7 @@ class GetDVCFileError(DvcException):
         )
 
 
-def get(url, path, out=None, rev=None):
+def get(url, path, out=None, rev=None, jobs=None):
     import shortuuid
 
     from dvc.dvcfile import is_valid_filename
@@ -52,6 +52,8 @@ def get(url, path, out=None, rev=None):
         ) as repo:
             from_info = PathInfo(repo.root_dir) / path
             to_info = PathInfo(out)
-            repo.repo_tree.download(from_info, to_info, follow_subrepos=False)
+            repo.repo_tree.download(
+                from_info, to_info, jobs=jobs, follow_subrepos=False
+            )
     finally:
         remove(tmp_dir)
diff --git a/dvc/repo/get_url.py b/dvc/repo/get_url.py
index 48475a61f..abc680d97 100644
--- a/dvc/repo/get_url.py
+++ b/dvc/repo/get_url.py
@@ -1,7 +1,7 @@
 import os
 
 
-def get_url(url, out=None):
+def get_url(url, out=None, jobs=None):
     import dvc.dependency as dependency
     import dvc.output as output
     from dvc.utils import resolve_output
@@ -17,4 +17,4 @@ def get_url(url, out=None):
     (out,) = output.loads_from(None, [out], use_cache=False)
     if not dep.exists:
         raise dep.DoesNotExistError(dep)
-    dep.download(out)
+    dep.download(out, jobs=jobs)
diff --git a/tests/unit/command/test_get.py b/tests/unit/command/test_get.py
index 6ac5e4bb5..db76ee2b0 100644
--- a/tests/unit/command/test_get.py
+++ b/tests/unit/command/test_get.py
@@ -4,7 +4,17 @@ from dvc.command.get import CmdGet
 
 def test_get(mocker):
     cli_args = parse_args(
-        ["get", "repo_url", "src", "--out", "out", "--rev", "version"]
+        [
+            "get",
+            "repo_url",
+            "src",
+            "--out",
+            "out",
+            "--rev",
+            "version",
+            "--jobs",
+            "4",
+        ]
     )
     assert cli_args.func == CmdGet
 
@@ -13,7 +23,9 @@ def test_get(mocker):
 
     assert cmd.run() == 0
 
-    m.assert_called_once_with("repo_url", path="src", out="out", rev="version")
+    m.assert_called_once_with(
+        "repo_url", path="src", out="out", rev="version", jobs=4
+    )
 
 
 def test_get_url(mocker, caplog):
diff --git a/tests/unit/command/test_get_url.py b/tests/unit/command/test_get_url.py
index 9b4a5d54c..71f61d8e8 100644
--- a/tests/unit/command/test_get_url.py
+++ b/tests/unit/command/test_get_url.py
@@ -3,7 +3,7 @@ from dvc.command.get_url import CmdGetUrl
 
 
 def test_get_url(mocker):
-    cli_args = parse_args(["get-url", "src", "out"])
+    cli_args = parse_args(["get-url", "src", "out", "-j", "5"])
     assert cli_args.func == CmdGetUrl
 
     cmd = cli_args.func(cli_args)
@@ -11,4 +11,4 @@ def test_get_url(mocker):
 
     assert cmd.run() == 0
 
-    m.assert_called_once_with("src", out="out")
+    m.assert_called_once_with("src", out="out", jobs=5)