From 023dec4aa6d11d80338803ba5200f4cd9cd62da7 Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <kupruser@gmail.com>
Date: Sat, 15 Aug 2020 20:35:39 +0300
Subject: [PATCH] tree: make get_hash return type and hash pair (#4397)

Currently we kinda assume that whatever is returned by `get_file_hash`
is of type self.PARAM_CHECKSUM, which is not actually true. E.g. for
http it might return `etag` or `md5`, but we don't distinguish between
those and call both `etag`. This is becoming more relevant for dir
hashes that are computed a few different ways (e.g. in-memory md5 or
upload to remote and get etag for the dir file).

Prerequisite for #4144 and #3069
---
 dvc/cache/base.py               |  8 +++++---
 dvc/cache/local.py              |  4 +++-
 dvc/dependency/repo.py          |  4 +++-
 dvc/output/base.py              |  2 +-
 dvc/repo/diff.py                |  2 +-
 dvc/repo/tree.py                |  9 ++++++---
 dvc/tree/azure.py               |  2 +-
 dvc/tree/base.py                | 28 +++++++++++++++-------------
 dvc/tree/gs.py                  |  7 +++++--
 dvc/tree/hdfs.py                |  2 +-
 dvc/tree/http.py                |  2 +-
 dvc/tree/local.py               |  2 +-
 dvc/tree/s3.py                  |  5 ++++-
 dvc/tree/ssh/__init__.py        |  2 +-
 dvc/tree/webdav.py              |  2 +-
 tests/func/test_tree.py         |  2 +-
 tests/unit/remote/test_azure.py |  2 +-
 17 files changed, 51 insertions(+), 34 deletions(-)

diff --git a/dvc/cache/base.py b/dvc/cache/base.py
index a6c5596d7..422a2548a 100644
--- a/dvc/cache/base.py
+++ b/dvc/cache/base.py
@@ -127,7 +127,9 @@ class CloudCache:
             logger.debug("cache for '%s'('%s') has changed.", path_info, hash_)
             return True
 
-        actual = self.tree.get_hash(path_info)
+        typ, actual = self.tree.get_hash(path_info)
+        assert typ == self.tree.PARAM_CHECKSUM
+
         if hash_ != actual:
             logger.debug(
                 "hash value '%s' for '%s' has changed (actual '%s').",
@@ -312,7 +314,7 @@ class CloudCache:
             )
             return False
 
-        actual = self.tree.get_hash(cache_info)
+        _, actual = self.tree.get_hash(cache_info)
 
         logger.debug(
             "cache '%s' expected '%s' actual '%s'", cache_info, hash_, actual,
@@ -358,7 +360,7 @@ class CloudCache:
         return self.changed_cache_file(hash_)
 
     def already_cached(self, path_info):
-        current = self.tree.get_hash(path_info)
+        _, current = self.tree.get_hash(path_info)
 
         if not current:
             return False
diff --git a/dvc/cache/local.py b/dvc/cache/local.py
index ec4d1d267..19e51cb1a 100644
--- a/dvc/cache/local.py
+++ b/dvc/cache/local.py
@@ -91,7 +91,9 @@ class LocalCache(CloudCache):
     def already_cached(self, path_info):
         assert path_info.scheme in ["", "local"]
 
-        current_md5 = self.tree.get_hash(path_info)
+        typ, current_md5 = self.tree.get_hash(path_info)
+
+        assert typ == "md5"
 
         if not current_md5:
             return False
diff --git a/dvc/dependency/repo.py b/dvc/dependency/repo.py
index 7dffd2eca..346aaa925 100644
--- a/dvc/dependency/repo.py
+++ b/dvc/dependency/repo.py
@@ -64,7 +64,9 @@ class RepoDependency(LocalDependency):
 
                 # We are polluting our repo cache with some dir listing here
                 if tree.isdir(path):
-                    return self.repo.cache.local.tree.get_hash(path, tree=tree)
+                    return self.repo.cache.local.tree.get_hash(
+                        path, tree=tree
+                    )[1]
                 return tree.get_file_hash(path)
 
     def workspace_status(self):
diff --git a/dvc/output/base.py b/dvc/output/base.py
index bdaded877..54f30b54a 100644
--- a/dvc/output/base.py
+++ b/dvc/output/base.py
@@ -184,7 +184,7 @@ class BaseOutput:
         self.info[self.tree.PARAM_CHECKSUM] = checksum
 
     def get_checksum(self):
-        return self.tree.get_hash(self.path_info)
+        return self.tree.get_hash(self.path_info)[1]
 
     @property
     def is_dir_checksum(self):
diff --git a/dvc/repo/diff.py b/dvc/repo/diff.py
index a9decaefc..19838aa77 100644
--- a/dvc/repo/diff.py
+++ b/dvc/repo/diff.py
@@ -37,7 +37,7 @@ def diff(self, a_rev="HEAD", b_rev=None):
 
         def _to_checksum(output):
             if on_working_tree:
-                return self.cache.local.tree.get_hash(output.path_info)
+                return self.cache.local.tree.get_hash(output.path_info)[1]
             return output.checksum
 
         def _exists(output):
diff --git a/dvc/repo/tree.py b/dvc/repo/tree.py
index 9140c6fa7..730a7aa50 100644
--- a/dvc/repo/tree.py
+++ b/dvc/repo/tree.py
@@ -238,8 +238,11 @@ class DvcTree(BaseTree):  # pylint:disable=abstract-method
             raise OutputNotFoundError
         out = outs[0]
         if out.is_dir_checksum:
-            return self._get_granular_checksum(path_info, out)
-        return out.checksum
+            return (
+                out.tree.PARAM_CHECKSUM,
+                self._get_granular_checksum(path_info, out),
+            )
+        return out.tree.PARAM_CHECKSUM, out.checksum
 
 
 class RepoTree(BaseTree):  # pylint:disable=abstract-method
@@ -504,7 +507,7 @@ class RepoTree(BaseTree):  # pylint:disable=abstract-method
                 return dvc_tree.get_file_hash(path_info)
             except OutputNotFoundError:
                 pass
-        return file_md5(path_info, self)[0]
+        return self.PARAM_CHECKSUM, file_md5(path_info, self)[0]
 
     def copytree(self, top, dest):
         top = PathInfo(top)
diff --git a/dvc/tree/azure.py b/dvc/tree/azure.py
index 78ada11b8..5ef3f2594 100644
--- a/dvc/tree/azure.py
+++ b/dvc/tree/azure.py
@@ -153,7 +153,7 @@ class AzureTree(BaseTree):
         ).delete_blob()
 
     def get_file_hash(self, path_info):
-        return self.get_etag(path_info)
+        return self.PARAM_CHECKSUM, self.get_etag(path_info)
 
     def _upload(
         self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
diff --git a/dvc/tree/base.py b/dvc/tree/base.py
index 1860bafab..3494e3272 100644
--- a/dvc/tree/base.py
+++ b/dvc/tree/base.py
@@ -242,7 +242,7 @@ class BaseTree:
         )
 
         if not self.exists(path_info):
-            return None
+            return self.PARAM_CHECKSUM, None
 
         # pylint: disable=assignment-from-none
         hash_ = self.state.get(path_info)
@@ -260,17 +260,17 @@ class BaseTree:
             hash_ = None
 
         if hash_:
-            return hash_
+            return self.PARAM_CHECKSUM, hash_
 
         if self.isdir(path_info):
-            hash_ = self.get_dir_hash(path_info, **kwargs)
+            typ, hash_ = self.get_dir_hash(path_info, **kwargs)
         else:
-            hash_ = self.get_file_hash(path_info)
+            typ, hash_ = self.get_file_hash(path_info)
 
         if hash_ and self.exists(path_info):
             self.state.save(path_info, hash_)
 
-        return hash_
+        return typ, hash_
 
     def get_file_hash(self, path_info):
         raise NotImplementedError
@@ -294,7 +294,8 @@ class BaseTree:
         return "".join(parts)
 
     def save_info(self, path_info, **kwargs):
-        return {self.PARAM_CHECKSUM: self.get_hash(path_info, **kwargs)}
+        typ, hash_ = self.get_hash(path_info, **kwargs)
+        return {typ: hash_}
 
     def _calculate_hashes(self, file_infos):
         file_infos = list(file_infos)
@@ -305,9 +306,10 @@ class BaseTree:
         ) as pbar:
             worker = pbar.wrap_fn(self.get_file_hash)
             with ThreadPoolExecutor(max_workers=self.hash_jobs) as executor:
-                tasks = executor.map(worker, file_infos)
-                hashes = dict(zip(file_infos, tasks))
-        return hashes
+                hashes = (
+                    value for typ, value in executor.map(worker, file_infos)
+                )
+                return dict(zip(file_infos, hashes))
 
     def _collect_dir(self, path_info, **kwargs):
         file_infos = set()
@@ -344,7 +346,7 @@ class BaseTree:
         return sorted(result, key=itemgetter(self.PARAM_RELPATH))
 
     def _save_dir_info(self, dir_info):
-        hash_, tmp_info = self._get_dir_info_hash(dir_info)
+        typ, hash_, tmp_info = self._get_dir_info_hash(dir_info)
         new_info = self.cache.tree.hash_to_path_info(hash_)
         if self.cache.changed_cache_file(hash_):
             self.cache.tree.makedirs(new_info.parent)
@@ -354,7 +356,7 @@ class BaseTree:
 
         self.state.save(new_info, hash_)
 
-        return hash_
+        return typ, hash_
 
     def _get_dir_info_hash(self, dir_info):
         tmp = tempfile.NamedTemporaryFile(delete=False).name
@@ -366,8 +368,8 @@ class BaseTree:
         to_info = tree.path_info / tmp_fname("")
         tree.upload(from_info, to_info, no_progress_bar=True)
 
-        hash_ = tree.get_file_hash(to_info) + self.CHECKSUM_DIR_SUFFIX
-        return hash_, to_info
+        typ, hash_ = tree.get_file_hash(to_info)
+        return typ, hash_ + self.CHECKSUM_DIR_SUFFIX, to_info
 
     def upload(self, from_info, to_info, name=None, no_progress_bar=False):
         if not hasattr(self, "_upload"):
diff --git a/dvc/tree/gs.py b/dvc/tree/gs.py
index 767154de3..c7055a73a 100644
--- a/dvc/tree/gs.py
+++ b/dvc/tree/gs.py
@@ -190,11 +190,14 @@ class GSTree(BaseTree):
         path = path_info.path
         blob = self.gs.bucket(bucket).get_blob(path)
         if not blob:
-            return None
+            return self.PARAM_CHECKSUM, None
 
         b64_md5 = blob.md5_hash
         md5 = base64.b64decode(b64_md5)
-        return codecs.getencoder("hex")(md5)[0].decode("utf-8")
+        return (
+            self.PARAM_CHECKSUM,
+            codecs.getencoder("hex")(md5)[0].decode("utf-8"),
+        )
 
     def _upload(self, from_file, to_info, name=None, no_progress_bar=False):
         bucket = self.gs.bucket(to_info.bucket)
diff --git a/dvc/tree/hdfs.py b/dvc/tree/hdfs.py
index cdc9a2c33..be2f48a6e 100644
--- a/dvc/tree/hdfs.py
+++ b/dvc/tree/hdfs.py
@@ -167,7 +167,7 @@ class HDFSTree(BaseTree):
         stdout = self.hadoop_fs(
             f"checksum {path_info.url}", user=path_info.user
         )
-        return self._group(regex, stdout, "checksum")
+        return self.PARAM_CHECKSUM, self._group(regex, stdout, "checksum")
 
     def _upload(self, from_file, to_info, **_kwargs):
         with self.hdfs(to_info) as hdfs:
diff --git a/dvc/tree/http.py b/dvc/tree/http.py
index dfe6f8bb9..e530c119d 100644
--- a/dvc/tree/http.py
+++ b/dvc/tree/http.py
@@ -136,7 +136,7 @@ class HTTPTree(BaseTree):  # pylint:disable=abstract-method
                 "Content-MD5 header for '{url}'".format(url=url)
             )
 
-        return etag
+        return self.PARAM_CHECKSUM, etag
 
     def _download(self, from_info, to_file, name=None, no_progress_bar=False):
         response = self.request("GET", from_info.url, stream=True)
diff --git a/dvc/tree/local.py b/dvc/tree/local.py
index 994277aea..6eaad240a 100644
--- a/dvc/tree/local.py
+++ b/dvc/tree/local.py
@@ -309,7 +309,7 @@ class LocalTree(BaseTree):
         return stat.S_IMODE(mode) == self.CACHE_MODE
 
     def get_file_hash(self, path_info):
-        return file_md5(path_info)[0]
+        return self.PARAM_CHECKSUM, file_md5(path_info)[0]
 
     @staticmethod
     def getsize(path_info):
diff --git a/dvc/tree/s3.py b/dvc/tree/s3.py
index 44368d746..47e5e19db 100644
--- a/dvc/tree/s3.py
+++ b/dvc/tree/s3.py
@@ -318,7 +318,10 @@ class S3Tree(BaseTree):
             raise ETagMismatchError(etag, cached_etag)
 
     def get_file_hash(self, path_info):
-        return self.get_etag(self.s3, path_info.bucket, path_info.path)
+        return (
+            self.PARAM_CHECKSUM,
+            self.get_etag(self.s3, path_info.bucket, path_info.path),
+        )
 
     def _upload(self, from_file, to_info, name=None, no_progress_bar=False):
         total = os.path.getsize(from_file)
diff --git a/dvc/tree/ssh/__init__.py b/dvc/tree/ssh/__init__.py
index 26bee1938..c5ecccbaa 100644
--- a/dvc/tree/ssh/__init__.py
+++ b/dvc/tree/ssh/__init__.py
@@ -238,7 +238,7 @@ class SSHTree(BaseTree):
             raise NotImplementedError
 
         with self.ssh(path_info) as ssh:
-            return ssh.md5(path_info.path)
+            return self.PARAM_CHECKSUM, ssh.md5(path_info.path)
 
     def getsize(self, path_info):
         with self.ssh(path_info) as ssh:
diff --git a/dvc/tree/webdav.py b/dvc/tree/webdav.py
index 48c48b89a..8909a7754 100644
--- a/dvc/tree/webdav.py
+++ b/dvc/tree/webdav.py
@@ -142,7 +142,7 @@ class WebDAVTree(BaseTree):  # pylint:disable=abstract-method
                 "Content-MD5 header for '{url}'".format(url=path_info.url)
             )
 
-        return etag
+        return self.PARAM_CHECKSUM, etag
 
     # Checks whether path points to directory
     def isdir(self, path_info):
diff --git a/tests/func/test_tree.py b/tests/func/test_tree.py
index 7f90196be..aec286243 100644
--- a/tests/func/test_tree.py
+++ b/tests/func/test_tree.py
@@ -211,7 +211,7 @@ def test_repotree_cache_save(tmp_dir, dvc, scm, erepo_dir, local_cloud):
     # into dvc.cache, not fetched or streamed from a remote
     tree = RepoTree(erepo_dir.dvc, stream=True)
     expected = [
-        tree.get_file_hash(PathInfo(erepo_dir / path))
+        tree.get_file_hash(PathInfo(erepo_dir / path))[1]
         for path in ("dir/bar", "dir/subdir/foo")
     ]
 
diff --git a/tests/unit/remote/test_azure.py b/tests/unit/remote/test_azure.py
index a0614bcf6..0f49045e3 100644
--- a/tests/unit/remote/test_azure.py
+++ b/tests/unit/remote/test_azure.py
@@ -36,7 +36,7 @@ def test_get_file_hash(tmp_dir, azure):
     to_info = azure
     tree.upload(PathInfo("foo"), to_info)
     assert tree.exists(to_info)
-    hash_ = tree.get_file_hash(to_info)
+    _, hash_ = tree.get_file_hash(to_info)
     assert hash_
     assert isinstance(hash_, str)
     assert hash_.strip("'").strip('"') == hash_