_artifactcache/cascache.py: Add _ensure_blob helper

This adds directory objects to the local repository before downloading
files in the directory. However, artifact references are still stored
only after downloading the complete directory and thus, there won't be
dangling references. This will anyway be required for partial download
support.

1 files changed, 33 insertions, 24 deletions

diff --git a/buildstream/_artifactcache/cascache.py b/buildstream/_artifactcache/cascache.py
index 8a0eb5195..619613a41 100644
--- a/buildstream/_artifactcache/cascache.py
+++ b/buildstream/_artifactcache/cascache.py
@@ -803,6 +803,31 @@ class CASCache(ArtifactCache):
         out.flush()
         assert digest.size_bytes == os.fstat(out.fileno()).st_size
 
+    # _ensure_blob():
+    #
+    # Fetch and add blob if it's not already local.
+    #
+    # Args:
+    #     remote (Remote): The remote to use.
+    #     digest (Digest): Digest object for the blob to fetch.
+    #
+    # Returns:
+    #     (str): The path of the object
+    #
+    def _ensure_blob(self, remote, digest):
+        objpath = self.objpath(digest)
+        if os.path.exists(objpath):
+            # already in local repository
+            return objpath
+
+        with tempfile.NamedTemporaryFile(dir=self.tmpdir) as f:
+            self._fetch_blob(remote, digest, f)
+
+            added_digest = self.add_object(path=f.name)
+            assert added_digest.hash == digest.hash
+
+        return objpath
+
     # _fetch_directory():
     #
     # Fetches remote directory and adds it to content addressable store.
@@ -821,34 +846,18 @@ class CASCache(ArtifactCache):
             # already in local cache
             return
 
-        with tempfile.NamedTemporaryFile(dir=self.tmpdir) as out:
-            self._fetch_blob(remote, dir_digest, out)
-
-            directory = remote_execution_pb2.Directory()
-
-            with open(out.name, 'rb') as f:
-                directory.ParseFromString(f.read())
-
-            for filenode in directory.files:
-                fileobjpath = self.objpath(filenode.digest)
-                if os.path.exists(fileobjpath):
-                    # already in local cache
-                    continue
+        objpath = self._ensure_blob(remote, dir_digest)
 
-                with tempfile.NamedTemporaryFile(dir=self.tmpdir) as f:
-                    self._fetch_blob(remote, filenode.digest, f)
+        directory = remote_execution_pb2.Directory()
 
-                    digest = self.add_object(path=f.name)
-                    assert digest.hash == filenode.digest.hash
+        with open(objpath, 'rb') as f:
+            directory.ParseFromString(f.read())
 
-            for dirnode in directory.directories:
-                self._fetch_directory(remote, dirnode.digest)
+        for filenode in directory.files:
+            self._ensure_blob(remote, filenode.digest)
 
-            # Place directory blob only in final location when we've
-            # downloaded all referenced blobs to avoid dangling
-            # references in the repository.
-            digest = self.add_object(path=out.name)
-            assert digest.hash == dir_digest.hash
+        for dirnode in directory.directories:
+            self._fetch_directory(remote, dirnode.digest)
 
 
 # Represents a single remote CAS cache.