cascache.py: Also fetch file blobs in _fetch_directory()

This eliminates code duplication in `ArtifactCache`, `SourceCache` and
`ElementSourcesCache`.

4 files changed, 7 insertions, 22 deletions

diff --git a/src/buildstream/_artifactcache.py b/src/buildstream/_artifactcache.py
index 09804fe01..efced2807 100644
--- a/src/buildstream/_artifactcache.py
+++ b/src/buildstream/_artifactcache.py
@@ -504,13 +504,6 @@ class ArtifactCache(AssetCache):
     #    blobs not existing on the server.
     #
     def _pull_artifact_storage(self, element, key, artifact_digest, remote, pull_buildtrees=False):
-        def __pull_digest(digest):
-            self.cas._fetch_directory(remote, digest)
-            required_blobs = self.cas.required_blobs_for_directory(digest)
-            missing_blobs = self.cas.local_missing_blobs(required_blobs)
-            if missing_blobs:
-                self.cas.fetch_blobs(remote, missing_blobs)
-
         artifact_name = element.get_artifact_name(key=key)
 
         try:
@@ -527,10 +520,10 @@ class ArtifactCache(AssetCache):
                 f.write(artifact.SerializeToString())
 
             if str(artifact.files):
-                __pull_digest(artifact.files)
+                self.cas._fetch_directory(remote, artifact.files)
 
             if pull_buildtrees and str(artifact.buildtree):
-                __pull_digest(artifact.buildtree)
+                self.cas._fetch_directory(remote, artifact.buildtree)
 
             digests = [artifact.low_diversity_meta, artifact.high_diversity_meta]
             if str(artifact.public_data):
diff --git a/src/buildstream/_cas/cascache.py b/src/buildstream/_cas/cascache.py
index 42e2244c5..d41d6ad71 100644
--- a/src/buildstream/_cas/cascache.py
+++ b/src/buildstream/_cas/cascache.py
@@ -509,8 +509,7 @@ class CASCache:
     #
     # Fetches remote directory and adds it to content addressable store.
     #
-    # This recursively fetches directory objects but doesn't fetch any
-    # files.
+    # This recursively fetches directory objects and files.
     #
     # Args:
     #     remote (Remote): The remote to use.
@@ -536,6 +535,9 @@ class CASCache:
                 "Failed to fetch directory tree {}: {}: {}".format(dir_digest.hash, e.code().name, e.details())
             ) from e
 
+        required_blobs = self.required_blobs_for_directory(dir_digest)
+        self.fetch_blobs(remote, required_blobs)
+
     def _fetch_tree(self, remote, digest):
         objpath = self._ensure_blob(remote, digest)
 
diff --git a/src/buildstream/_elementsourcescache.py b/src/buildstream/_elementsourcescache.py
index 84e7633e5..9a617b8d6 100644
--- a/src/buildstream/_elementsourcescache.py
+++ b/src/buildstream/_elementsourcescache.py
@@ -295,13 +295,6 @@ class ElementSourcesCache(AssetCache):
     #    blobs not existing on the server.
     #
     def _pull_source_storage(self, key, source_digest, remote):
-        def __pull_digest(digest):
-            self.cas._fetch_directory(remote, digest)
-            required_blobs = self.cas.required_blobs_for_directory(digest)
-            missing_blobs = self.cas.local_missing_blobs(required_blobs)
-            if missing_blobs:
-                self.cas.fetch_blobs(remote, missing_blobs)
-
         try:
             # Fetch and parse source proto
             self.cas.fetch_blobs(remote, [source_digest])
@@ -314,7 +307,7 @@ class ElementSourcesCache(AssetCache):
             with utils.save_file_atomic(source_path, mode="wb") as f:
                 f.write(source.SerializeToString())
 
-            __pull_digest(source.files)
+            self.cas._fetch_directory(remote, source.files)
         except grpc.RpcError as e:
             if e.code() != grpc.StatusCode.NOT_FOUND:
                 raise SourceCacheError("Failed to pull source with status {}: {}".format(e.code().name, e.details()))
diff --git a/src/buildstream/_sourcecache.py b/src/buildstream/_sourcecache.py
index 76c22efbd..37d990b4f 100644
--- a/src/buildstream/_sourcecache.py
+++ b/src/buildstream/_sourcecache.py
@@ -146,9 +146,6 @@ class SourceCache(AssetCache):
 
                 # Fetch source blobs
                 self.cas._fetch_directory(remote, source_digest)
-                required_blobs = self.cas.required_blobs_for_directory(source_digest)
-                missing_blobs = self.cas.local_missing_blobs(required_blobs)
-                self.cas.fetch_blobs(remote, missing_blobs)
 
                 source.info("Pulled source {} <- {}".format(display_key, remote))
                 return True