From 7806afd4ad46a0742fc928dd2ffbfc476ad27b27 Mon Sep 17 00:00:00 2001
From: Peter Rowlands <peter@pmrowla.com>
Date: Mon, 24 Jul 2023 15:23:21 +0900
Subject: [PATCH] fs: prefer generic.copy over fsspec fs.get()

---
 dvc/dependency/base.py         |  2 +-
 dvc/fs/__init__.py             | 34 +++++++++++++++++++++++++++++++---
 dvc/repo/get.py                | 19 +++++++------------
 dvc/repo/get_url.py            |  2 +-
 tests/func/repro/test_repro.py |  6 ++++--
 5 files changed, 44 insertions(+), 19 deletions(-)

diff --git a/dvc/dependency/base.py b/dvc/dependency/base.py
index 2aa90a549..cc94f610d 100644
--- a/dvc/dependency/base.py
+++ b/dvc/dependency/base.py
@@ -50,7 +50,7 @@ class Dependency(Output):
             self.fs_path = self.fs.path.version_path(self.fs_path, self.meta.version_id)
 
     def download(self, to, jobs=None):
-        fs_download(self.fs, self.fs_path, to, jobs=jobs)
+        fs_download(self.fs, self.fs_path, to.fs_path, jobs=jobs)
 
     def save(self):
         super().save()
diff --git a/dvc/fs/__init__.py b/dvc/fs/__init__.py
index 59ad0e5a1..8653d0fd4 100644
--- a/dvc/fs/__init__.py
+++ b/dvc/fs/__init__.py
@@ -1,3 +1,4 @@
+from typing import Optional
 from urllib.parse import urlparse
 
 from dvc_http import HTTPFileSystem, HTTPSFileSystem  # noqa: F401
@@ -18,7 +19,7 @@ from dvc_objects.fs import (  # noqa: F401
     system,
     utils,
 )
-from dvc_objects.fs.base import AnyFSPath, FileSystem  # noqa: F401
+from dvc_objects.fs.base import AnyFSPath, FileSystem  # noqa: F401, TCH001
 from dvc_objects.fs.errors import (  # noqa: F401
     AuthError,
     ConfigError,
@@ -48,12 +49,39 @@ known_implementations.update(
 # pylint: enable=unused-import
 
 
-def download(fs, fs_path, to, jobs=None):
+def download(fs: "FileSystem", fs_path: str, to: str, jobs: Optional[int] = None):
     with Callback.as_tqdm_callback(
         desc=f"Downloading {fs.path.name(fs_path)}",
         unit="files",
     ) as cb:
-        fs.get(fs_path, to.fs_path, batch_size=jobs, callback=cb)
+        # NOTE: We use dvc-objects generic.copy over fs.get since it makes file
+        # download atomic and avoids fsspec glob/regex path expansion.
+        if fs.isdir(fs_path):
+            from_infos = [
+                path
+                for path in fs.find(fs_path)
+                if not path.endswith(fs.path.flavour.sep)
+            ]
+            if not from_infos:
+                return localfs.makedirs(to, exist_ok=True)
+            to_infos = [
+                localfs.path.join(to, *fs.path.relparts(info, fs_path))
+                for info in from_infos
+            ]
+        else:
+            from_infos = [fs_path]
+            to_infos = [to]
+
+        cb.set_size(len(from_infos))
+        jobs = jobs or fs.jobs
+        generic.copy(
+            fs,
+            from_infos,
+            localfs,
+            to_infos,
+            callback=cb,
+            batch_size=jobs,
+        )
 
 
 def parse_external_url(url, config=None):
diff --git a/dvc/repo/get.py b/dvc/repo/get.py
index 178a46f59..239a08318 100644
--- a/dvc/repo/get.py
+++ b/dvc/repo/get.py
@@ -22,7 +22,6 @@ class GetDVCFileError(DvcException):
 def get(url, path, out=None, rev=None, jobs=None, force=False, config=None):
     from dvc.config import Config
     from dvc.dvcfile import is_valid_filename
-    from dvc.fs.callbacks import Callback
     from dvc.repo import Repo
 
     out = resolve_output(path, out, force=force)
@@ -40,6 +39,7 @@ def get(url, path, out=None, rev=None, jobs=None, force=False, config=None):
         uninitialized=True,
         config=config,
     ) as repo:
+        from dvc.fs import download
         from dvc.fs.data import DataFileSystem
 
         fs: Union[DataFileSystem, "DVCFileSystem"]
@@ -49,14 +49,9 @@ def get(url, path, out=None, rev=None, jobs=None, force=False, config=None):
         else:
             fs = repo.dvcfs
             fs_path = fs.from_os_path(path)
-
-        with Callback.as_tqdm_callback(
-            desc=f"Downloading {fs.path.name(path)}",
-            unit="files",
-        ) as cb:
-            fs.get(
-                fs_path,
-                os.path.abspath(out),
-                batch_size=jobs,
-                callback=cb,
-            )
+        download(
+            fs,
+            fs_path,
+            os.path.abspath(out),
+            jobs=jobs,
+        )
diff --git a/dvc/repo/get_url.py b/dvc/repo/get_url.py
index 22acb5f0c..7b6ab4f8f 100644
--- a/dvc/repo/get_url.py
+++ b/dvc/repo/get_url.py
@@ -14,4 +14,4 @@ def get_url(url, out=None, *, config=None, jobs=None, force=False):
     src_fs, src_path = parse_external_url(url, config)
     if not src_fs.exists(src_path):
         raise URLMissingError(url)
-    download(src_fs, src_path, out, jobs=jobs)
+    download(src_fs, src_path, out.fs_path, jobs=jobs)
diff --git a/tests/func/repro/test_repro.py b/tests/func/repro/test_repro.py
index ee4563a3b..d31fd2343 100644
--- a/tests/func/repro/test_repro.py
+++ b/tests/func/repro/test_repro.py
@@ -10,7 +10,7 @@ from funcy import lsplit
 from dvc.cli import main
 from dvc.dvcfile import LOCK_FILE, PROJECT_FILE
 from dvc.exceptions import CyclicGraphError, ReproductionError
-from dvc.fs import LocalFileSystem, system
+from dvc.fs import system
 from dvc.output import Output
 from dvc.stage import PipelineStage, Stage
 from dvc.stage.cache import RunCacheNotSupported
@@ -624,12 +624,14 @@ class TestReproAlreadyCached:
         assert stage.outs[0].hash_info != saved_stage.outs[0].hash_info
 
     def test_force_import(self, mocker, tmp_dir, dvc):
+        from dvc.dependency import base
+
         tmp_dir.dvc_gen("foo", "foo")
 
         ret = main(["import-url", "foo", "bar"])
         assert ret == 0
 
-        spy_get = mocker.spy(LocalFileSystem, "get")
+        spy_get = mocker.spy(base, "fs_download")
         spy_checkout = mocker.spy(Output, "checkout")
 
         assert main(["unfreeze", "bar.dvc"]) == 0