Reimplement cquery --output=streamed_proto.

- The current state of the output formats of cquery has a few forms now: - (UNCHANGED) cquery --output=proto|jsonproto|textproto --proto:include_configurations: A single CqueryResult of the specified --output format. - (NEW) cquery --output=streamed_proto --proto:include_configurations: Multiple length-delimited CqueryResult protos each containing a single ConfiguredTarget or Configuration. - (UNCHANGED) cquery --output=proto|jsonproto|textproto --noproto:include_configurations: A single QueryResult containing the corresponding Target(s) in the specified format. - (NEW) cquery --output=streamed_proto --noproto:include_configurations: Multiple length-delimited QueryResult protos each containing a single Target. - Added tests to ensure that all 4 (proto, jsonproto, textproto, streamed_proto) output formats are equivalent. PiperOrigin-RevId: 546794430 Change-Id: Ia027f070d008bfb721d88ff928392ebd5e694cd1
bazelbuild · Jul 10, 2023 · aa88357 · aa88357
1 parent 7845aca
commit aa88357
Show file tree

Hide file tree

Showing 5 changed files with 260 additions and 39 deletions.
diff --git a/...in/java/com/google/devtools/build/lib/query2/cquery/ConfiguredTargetQueryEnvironment.java b/...in/java/com/google/devtools/build/lib/query2/cquery/ConfiguredTargetQueryEnvironment.java
@@ -231,6 +231,15 @@ private static ImmutableMap<String, BuildConfigurationValue> getTransitiveConfig
             aspectResolver,
             OutputType.BINARY,
             ruleClassProvider),
+        new ProtoOutputFormatterCallback(
+            eventHandler,
+            cqueryOptions,
+            out,
+            skyframeExecutor,
+            accessor,
+            aspectResolver,
+            OutputType.DELIMITED_BINARY,
+            ruleClassProvider),
         new ProtoOutputFormatterCallback(
             eventHandler,
             cqueryOptions,

diff --git a/src/main/java/com/google/devtools/build/lib/query2/cquery/CqueryOptions.java b/src/main/java/com/google/devtools/build/lib/query2/cquery/CqueryOptions.java
@@ -45,8 +45,9 @@ public enum Transitions {
       effectTags = {OptionEffectTag.TERMINAL_OUTPUT},
       help =
           "The format in which the cquery results should be printed. Allowed values for cquery "
-              + "are: label, label_kind, textproto, transitions, proto, jsonproto. If you select "
-              + "'transitions', you also have to specify the --transitions=(lite|full) option.")
+              + "are: label, label_kind, textproto, transitions, proto, streamed_proto, jsonproto. "
+              + "If you select 'transitions', you also have to specify the "
+              + "--transitions=(lite|full) option.")
   public String outputFormat;
 
   @Option(

diff --git a/src/main/java/com/google/devtools/build/lib/query2/cquery/ProtoOutputFormatterCallback.java b/src/main/java/com/google/devtools/build/lib/query2/cquery/ProtoOutputFormatterCallback.java
@@ -15,14 +15,15 @@
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 
-import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Ordering;
 import com.google.devtools.build.lib.actions.BuildConfigurationEvent;
 import com.google.devtools.build.lib.analysis.AnalysisProtosV2;
 import com.google.devtools.build.lib.analysis.AnalysisProtosV2.Configuration;
+import com.google.devtools.build.lib.analysis.AnalysisProtosV2.CqueryResult;
+import com.google.devtools.build.lib.analysis.AnalysisProtosV2.CqueryResultOrBuilder;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.config.BuildOptions;
 import com.google.devtools.build.lib.analysis.config.ConfigMatchingProvider;
@@ -39,6 +40,7 @@
 import com.google.devtools.build.lib.query2.cquery.CqueryTransitionResolver.EvaluateException;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment.TargetAccessor;
 import com.google.devtools.build.lib.query2.proto.proto2api.Build;
+import com.google.devtools.build.lib.query2.proto.proto2api.Build.QueryResult;
 import com.google.devtools.build.lib.query2.query.aspectresolvers.AspectResolver;
 import com.google.devtools.build.lib.query2.query.output.ProtoOutputFormatter;
 import com.google.devtools.build.lib.skyframe.BuildConfigurationKey;
@@ -60,6 +62,7 @@ class ProtoOutputFormatterCallback extends CqueryThreadsafeCallback {
   /** Defines the types of proto output this class can handle. */
   public enum OutputType {
     BINARY("proto"),
+    DELIMITED_BINARY("streamed_proto"),
     TEXT("textproto"),
     JSON("jsonproto");
 
@@ -102,15 +105,14 @@ public ImmutableList<Configuration> getConfigurations() {
     }
   }
 
+  private CqueryResult.Builder cqueryResultBuilder;
   private final OutputType outputType;
   private final AspectResolver resolver;
   private final SkyframeExecutor skyframeExecutor;
   private final ConfigurationCache configurationCache = new ConfigurationCache();
   private final JsonFormat.Printer jsonPrinter = JsonFormat.printer();
   private final RuleClassProvider ruleClassProvider;
 
-  private AnalysisProtosV2.CqueryResult.Builder protoResult;
-
   private final Map<Label, Target> partialResultMap;
   private ConfiguredTarget currentTarget;
 
@@ -133,40 +135,80 @@ public ImmutableList<Configuration> getConfigurations() {
 
   @Override
   public void start() {
-    protoResult = AnalysisProtosV2.CqueryResult.newBuilder();
+    cqueryResultBuilder = AnalysisProtosV2.CqueryResult.newBuilder();
+  }
+
+  private static QueryResult queryResultFromCqueryResult(CqueryResultOrBuilder cqueryResult) {
+    Build.QueryResult.Builder queryResult = Build.QueryResult.newBuilder();
+    cqueryResult.getResultsList().forEach(ct -> queryResult.addTarget(ct.getTarget()));
+    return queryResult.build();
   }
 
   @Override
   public void close(boolean failFast) throws IOException {
-    if (!failFast && printStream != null) {
-      if (options.protoIncludeConfigurations) {
-        writeData(getProtoResult());
-      } else {
-        // Documentation promises that setting this flag to false means we convert directly
-        // to the build.proto format. This is hard to test in integration testing due to the way
-        // proto output is turned readable (codex). So change the following code with caution.
-        Build.QueryResult.Builder queryResult = Build.QueryResult.newBuilder();
-        protoResult.getResultsList().forEach(ct -> queryResult.addTarget(ct.getTarget()));
-        writeData(queryResult.build());
-      }
-      printStream.flush();
+    if (failFast || printStream == null) {
+      return;
+    }
+
+    // There are a few cases that affect the shape of the output:
+    //   1. --output=proto|textproto|jsonproto --proto:include_configurations =>
+    //        Writes a single CqueryResult containing all the ConfiguredTarget(s) and
+    //        Configuration(s) in the specified output format.
+    //   2. --output=streamed_proto --proto:include_configurations =>
+    //        Writes multiple length delimited CqueryResult protos, each containing a single
+    //        ConfiguredTarget or Configuration.
+    //   3. --output=proto|textproto|jsonproto --noproto:include_configurations =>
+    //        Writes a single QueryResult containing all the corresponding Target(s) in the
+    //        specified output format.
+    //   4.--output=streamed_proto --noproto:include_configurations =>
+    //        Writes multiple length delimited QueryResult protos, each containing a single Target.
+    switch (outputType) {
+      case BINARY:
+      case TEXT:
+      case JSON:
+        // Only at the end, we write the entire CqueryResult / QueryResult is written all together.
+        if (options.protoIncludeConfigurations) {
+          cqueryResultBuilder.addAllConfigurations(configurationCache.getConfigurations());
+        }
+        writeData(
+            options.protoIncludeConfigurations
+                ? cqueryResultBuilder.build()
+                : queryResultFromCqueryResult(cqueryResultBuilder));
+        break;
+      case DELIMITED_BINARY:
+        if (options.protoIncludeConfigurations) {
+          // The wrapped CqueryResult + ConfiguredTarget are already written in
+          // {@link #processOutput}, so we just need to write the Configuration(s) each wrapped in
+          // a CqueryResult.
+          for (Configuration configuration : configurationCache.getConfigurations()) {
+            writeData(
+                AnalysisProtosV2.CqueryResult.newBuilder()
+                    .addConfigurations(configuration)
+                    .build());
+          }
+        }
+        break;
     }
+
+    outputStream.flush();
+    printStream.flush();
   }
 
   private void writeData(Message message) throws IOException {
     switch (outputType) {
       case BINARY:
         message.writeTo(outputStream);
         break;
+      case DELIMITED_BINARY:
+        message.writeDelimitedTo(outputStream);
+        break;
       case TEXT:
         TextFormat.printer().print(message, printStream);
         break;
       case JSON:
         jsonPrinter.appendTo(message, printStream);
         printStream.append('\n');
         break;
-      default:
-        throw new IllegalStateException("Unknown outputType " + outputType.formatName());
     }
   }
 
@@ -175,14 +217,9 @@ public String getName() {
     return outputType.formatName();
   }
 
-  @VisibleForTesting
-  public AnalysisProtosV2.CqueryResult getProtoResult() {
-    protoResult.addAllConfigurations(configurationCache.getConfigurations());
-    return protoResult.build();
-  }
-
   @Override
-  public void processOutput(Iterable<ConfiguredTarget> partialResult) throws InterruptedException {
+  public void processOutput(Iterable<ConfiguredTarget> partialResult)
+      throws InterruptedException, IOException {
     partialResult.forEach(
         kct -> partialResultMap.put(kct.getOriginalLabel(), accessor.getTarget(kct)));
 
@@ -239,6 +276,7 @@ public void processOutput(Iterable<ConfiguredTarget> partialResult) throws Inter
           throw new InterruptedException(e.getMessage());
         }
       }
+
       builder.setTarget(targetBuilder);
 
       if (options.protoIncludeConfigurations) {
@@ -259,7 +297,20 @@ public void processOutput(Iterable<ConfiguredTarget> partialResult) throws Inter
         }
       }
 
-      protoResult.addResults(builder.build());
+      if (outputType == OutputType.DELIMITED_BINARY) {
+        // If --proto:include_configurations, we wrap the single ConfiguredTarget in a CqueryResult.
+        // If --noproto:include_configurations, we wrap the single Target in a QueryResult.
+        // Then we write either result delimited to the stream.
+        writeData(
+            options.protoIncludeConfigurations
+                ? CqueryResult.newBuilder().addResults(builder).build()
+                : QueryResult.newBuilder().addTarget(builder.getTarget()).build());
+      } else {
+        // Except --output=streamed_proto, all other output types require they be wrapped in a
+        // CqueryResult or QueryResult. So we instead of writing straight to the stream, we
+        // aggregate the results in a CqueryResult.Builder before writing in {@link #close}.
+        cqueryResultBuilder.addResults(builder.build());
+      }
     }
   }
 

diff --git a/src/test/java/com/google/devtools/build/lib/query2/cquery/BUILD b/src/test/java/com/google/devtools/build/lib/query2/cquery/BUILD
@@ -143,6 +143,8 @@ java_test(
         "//third_party:guava",
         "//third_party:junit4",
         "//third_party:truth",
+        "//third_party/protobuf:protobuf_java",
+        "//third_party/protobuf:protobuf_java_util",
     ],
 )