temp_directory.py: Add output directory upload support

#61

buildgrid/_app/bots/temp_directory.py

@@ -21,7 +21,8 @@ from google.protobuf import any_pb2
 
 from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2, remote_execution_pb2_grpc
 from buildgrid._protos.google.bytestream import bytestream_pb2_grpc
-from buildgrid.utils import output_file_maker, write_fetch_directory, parse_to_pb2_from_fetch
+from buildgrid.utils import write_fetch_directory, parse_to_pb2_from_fetch
+from buildgrid.utils import output_file_maker, output_directory_maker
 
 
 def work_temp_directory(context, lease):
@@ -88,11 +89,24 @@ def work_temp_directory(context, lease):
             if not os.path.isfile(file_path):
                 continue
 
+            # OutputFile.path should be relative to the working direcory:
             output_file, update_request = output_file_maker(file_path, working_directory)
 
             action_result.output_files.extend([output_file])
             update_requests.requests.extend([update_request])
 
+        for output_path in command.output_directories:
+            directory_path = os.path.join(working_directory, output_path)
+            # Missing outputs should simply be omitted in ActionResult:
+            if not os.path.isdir(directory_path):
+                continue
+
+            # OutputDirectory.path should be relative to the working direcory:
+            output_directory, update_request = output_directory_maker(directory_path, working_directory)
+
+            action_result.output_directories.extend([output_directory])
+            update_requests.requests.extend(update_request)
+
         stub_cas = remote_execution_pb2_grpc.ContentAddressableStorageStub(context.cas_channel)
         stub_cas.BatchUpdateBlobs(update_requests)
 

buildgrid/utils.py

@@ -179,7 +179,7 @@ def directory_maker(directory_path):
 
             node = remote_execution_pb2.FileNode()
             node.name = directory_entry.name
-            node.digest.copyFrom(node_digest)
+            node.digest.CopyFrom(node_digest)
             node.is_executable = os.access(directory_entry.path, os.X_OK)
 
             node_request = remote_execution_pb2.BatchUpdateBlobsRequest.Request(digest=node_digest)
@@ -211,10 +211,14 @@ def directory_maker(directory_path):
 
             symlinks.append(node)
 
+    files.sort(key=attrgetter('name'))
+    directories.sort(key=attrgetter('name'))
+    symlinks.sort(key=attrgetter('name'))
+
     directory = remote_execution_pb2.Directory()
-    directory.files.extend(files.sort(key=attrgetter('name')))
-    directory.directories.extend(directories.sort(key=attrgetter('name')))
-    directory.symlinks.extend(symlinks.sort(key=attrgetter('name')))
+    directory.files.extend(files)
+    directory.directories.extend(directories)
+    directory.symlinks.extend(symlinks)
 
     directory_blob = directory.SerializeToString()
     directory_digest = create_digest(directory_blob)
@@ -298,7 +302,7 @@ def output_directory_maker(directory_path, working_path):
     _, update_requests = tree_maker(directory_path)
 
     output_directory = remote_execution_pb2.OutputDirectory()
-    output_directory.tree_digest = update_requests[-1].digest
+    output_directory.tree_digest.CopyFrom(update_requests[-1].digest)
     output_directory.path = os.path.relpath(directory_path, start=working_path)
 
     output_directory_blob = output_directory.SerializeToString()
@@ -332,7 +336,7 @@ def tree_maker(directory_path):
     directory, child_directories, update_requests = directory_maker(directory_path)
 
     tree = remote_execution_pb2.Tree()
-    tree.children.extend([child_directories])
+    tree.children.extend(child_directories)
     tree.root.CopyFrom(directory)
 
     tree_blob = tree.SerializeToString()