buildgrid/_app/commands/cmd_cas.py

@@ -147,33 +147,35 @@ def _create_digest(digest_string):
     return digest
 
 
-@cli.command('download-file', short_help="Download a file from the CAS server.")
-@click.argument('digest-string', nargs=1, type=click.STRING, required=True)
-@click.argument('file-path', nargs=1, type=click.Path(exists=False), required=True)
+@cli.command('download-file', short_help="Download one or more files from the CAS server. "
+                                         "(Specified as a space-separated list of DIGEST FILE_PATH)")
+@click.argument('digest-path-list', nargs=-1, type=str, required=True)  # 'digest path' pairs
 @click.option('--verify', is_flag=True, show_default=True,
               help="Check downloaded file's integrity.")
 @pass_context
-def download_file(context, digest_string, file_path, verify):
-    if os.path.exists(file_path):
-        click.echo("Error: Invalid value for " +
-                   "path=[{}] already exists.".format(file_path), err=True)
-        return
-
-    digest = _create_digest(digest_string)
-    with download(context.channel, instance=context.instance_name) as downloader:
-        downloader.download_file(digest, file_path)
-
-    if verify:
-        file_digest = create_digest(read_file(file_path))
-        if file_digest != digest:
-            click.echo("Error: Failed to verify path=[{}]".format(file_path), err=True)
-            return
-
-    if os.path.isfile(file_path):
-        click.echo("Success: Pulled path=[{}] from digest=[{}/{}]"
-                   .format(file_path, digest.hash, digest.size_bytes))
-    else:
-        click.echo('Error: Failed pulling "{}"'.format(file_path), err=True)
+def download_file(context, digest_path_list, verify):
+    for (digest_string, file_path) in zip(digest_path_list[0::2],
+                                          digest_path_list[1::2]):
+        if os.path.exists(file_path):
+            click.echo("Error: Invalid value for " +
+                       "path=[{}] already exists.".format(file_path), err=True)
+            continue
+
+        digest = _create_digest(digest_string)
+        with download(context.channel, instance=context.instance_name) as downloader:
+            downloader.download_file(digest, file_path)
+
+        if verify:
+            file_digest = create_digest(read_file(file_path))
+            if file_digest != digest:
+                click.echo("Error: Failed to verify path=[{}]".format(file_path), err=True)
+                continue
+
+        if os.path.isfile(file_path):
+            click.echo("Success: Pulled path=[{}] from digest=[{}/{}]"
+                       .format(file_path, digest.hash, digest.size_bytes))
+        else:
+            click.echo('Error: Failed pulling "{}"'.format(file_path), err=True)
 
 
 @cli.command('download-dir', short_help="Download a directory from the CAS server.")

buildgrid/server/cas/instance.py

@@ -24,6 +24,7 @@ import logging
 from buildgrid._exceptions import InvalidArgumentError, NotFoundError, OutOfRangeError
 from buildgrid._protos.google.bytestream import bytestream_pb2
 from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2 as re_pb2
+from buildgrid._protos.google.rpc import code_pb2, status_pb2
 from buildgrid.settings import HASH, HASH_LENGTH, MAX_REQUEST_SIZE, MAX_REQUEST_COUNT
 from buildgrid.utils import get_hash_type
 
@@ -70,6 +71,35 @@ class ContentAddressableStorageInstance:
 
         return response
 
+    def batch_read_blobs(self, digests):
+        storage = self._storage
+
+        response = re_pb2.BatchReadBlobsResponse()
+
+        requested_bytes = sum((digest.size_bytes for digest in digests))
+        max_batch_size = self.max_batch_total_size_bytes()
+
+        if requested_bytes > max_batch_size:
+            raise InvalidArgumentError('Combined total size of blobs exceeds '
+                                       'server limit. '
+                                       '({} > {} [byte])'.format(requested_bytes,
+                                                                 max_batch_size))
+
+        for digest in digests:
+            response_proto = response.responses.add()
+            response_proto.digest.CopyFrom(digest)
+
+            blob = storage.get_blob(digest)
+            if blob:
+                response_proto.data = blob.read()
+                status_code = code_pb2.OK
+            else:
+                status_code = code_pb2.NOT_FOUND
+
+            response_proto.status.CopyFrom(status_pb2.Status(code=status_code))
+
+        return response
+
     def get_tree(self, request):
         storage = self._storage
 

buildgrid/server/cas/service.py

@@ -86,8 +86,15 @@ class ContentAddressableStorageService(remote_execution_pb2_grpc.ContentAddressa
     def BatchReadBlobs(self, request, context):
         self.__logger.debug("BatchReadBlobs request from [%s]", context.peer())
 
-        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
-        context.set_details('Method not implemented!')
+        try:
+            instance = self._get_instance(request.instance_name)
+            response = instance.batch_read_blobs(request.digests)
+            return response
+
+        except InvalidArgumentError as e:
+            self.__logger.error(e)
+            context.set_details(str(e))
+            context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
 
         return remote_execution_pb2.BatchReadBlobsResponse()