casremote.py: Use FetchMissingBlobs in _fetch_blob()

2 files changed, 25 insertions, 32 deletions

diff --git a/src/buildstream/_cas/cascache.py b/src/buildstream/_cas/cascache.py
index b8d651a0c..5ebf54601 100644
--- a/src/buildstream/_cas/cascache.py
+++ b/src/buildstream/_cas/cascache.py
@@ -820,11 +820,7 @@ class CASCache():
             # already in local repository
             return objpath
 
-        with self._temporary_object() as f:
-            remote._fetch_blob(digest, f)
-
-            added_digest = self.add_object(path=f.name, link_directly=True)
-            assert added_digest.hash == digest.hash
+        remote._fetch_blob(digest)
 
         return objpath
 
@@ -915,20 +911,18 @@ class CASCache():
         self._fetch_directory_batch(remote, batch, fetch_queue, fetch_next_queue)
 
     def _fetch_tree(self, remote, digest):
-        # download but do not store the Tree object
-        with utils._tempnamedfile(dir=self.tmpdir) as out:
-            remote._fetch_blob(digest, out)
+        objpath = self._ensure_blob(remote, digest)
 
-            tree = remote_execution_pb2.Tree()
+        tree = remote_execution_pb2.Tree()
 
-            with open(out.name, 'rb') as f:
-                tree.ParseFromString(f.read())
+        with open(objpath, 'rb') as f:
+            tree.ParseFromString(f.read())
 
-            tree.children.extend([tree.root])
-            for directory in tree.children:
-                dirbuffer = directory.SerializeToString()
-                dirdigest = self.add_object(buffer=dirbuffer)
-                assert dirdigest.size_bytes == len(dirbuffer)
+        tree.children.extend([tree.root])
+        for directory in tree.children:
+            dirbuffer = directory.SerializeToString()
+            dirdigest = self.add_object(buffer=dirbuffer)
+            assert dirdigest.size_bytes == len(dirbuffer)
 
         return dirdigest
 
diff --git a/src/buildstream/_cas/casremote.py b/src/buildstream/_cas/casremote.py
index e1848cd20..2a6028bf8 100644
--- a/src/buildstream/_cas/casremote.py
+++ b/src/buildstream/_cas/casremote.py
@@ -262,22 +262,21 @@ class CASRemote():
     ################################################
     #             Local Private Methods            #
     ################################################
-    def _fetch_blob(self, digest, stream):
-        if self.instance_name:
-            resource_name = '/'.join([self.instance_name, 'blobs',
-                                      digest.hash, str(digest.size_bytes)])
-        else:
-            resource_name = '/'.join(['blobs',
-                                      digest.hash, str(digest.size_bytes)])
-
-        request = bytestream_pb2.ReadRequest()
-        request.resource_name = resource_name
-        request.read_offset = 0
-        for response in self.bytestream.Read(request):
-            stream.write(response.data)
-        stream.flush()
-
-        assert digest.size_bytes == os.fstat(stream.fileno()).st_size
+    def _fetch_blob(self, digest):
+        local_cas = self.cascache._get_local_cas()
+        request = local_cas_pb2.FetchMissingBlobsRequest()
+        request.instance_name = self.local_cas_instance_name
+        request_digest = request.blob_digests.add()
+        request_digest.CopyFrom(digest)
+        response = local_cas.FetchMissingBlobs(request)
+        for blob_response in response.responses:
+            if blob_response.status.code == code_pb2.NOT_FOUND:
+                raise BlobNotFound(response.digest.hash, "Failed to download blob {}: {}".format(
+                    blob_response.digest.hash, blob_response.status.code))
+
+            if blob_response.status.code != code_pb2.OK:
+                raise CASRemoteError("Failed to download blob {}: {}".format(
+                    blob_response.digest.hash, blob_response.status.code))
 
     def _send_blob(self, digest, stream, u_uid=uuid.uuid4()):
         if self.instance_name: