Introduce a new method on the StreamedFormatter interface for creating a callback for streaming a precomputed result.

--
MOS_MIGRATED_REVID=133720742

5 files changed, 65 insertions, 11 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/query2/output/OutputFormatter.java b/src/main/java/com/google/devtools/build/lib/query2/output/OutputFormatter.java
index f47d94942b..8a3c4b6318 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/output/OutputFormatter.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/output/OutputFormatter.java
@@ -31,6 +31,7 @@ import com.google.devtools.build.lib.packages.License;
 import com.google.devtools.build.lib.packages.Rule;
 import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.query2.engine.OutputFormatterCallback;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment;
 import com.google.devtools.build.lib.query2.output.QueryOptions.OrderOutput;
 import com.google.devtools.build.lib.syntax.EvalUtils;
 import com.google.devtools.build.lib.syntax.Printer;
@@ -154,7 +155,8 @@ public abstract class OutputFormatter implements Serializable {
    * the graph maintained by the QueryEnvironment), and print it to "out".
    */
   public abstract void output(QueryOptions options, Digraph<Target> result, PrintStream out,
-      AspectResolver aspectProvider) throws IOException, InterruptedException;
+      AspectResolver aspectProvider)
+          throws IOException, InterruptedException;
 
   /**
    * Unordered streamed output formatter (wrt. dependency ordering).
@@ -177,8 +179,18 @@ public abstract class OutputFormatter implements Serializable {
      *
      * <p>Takes any options specified via the most recent call to {@link #setOptions} into
      * consideration.
+     *
+     * <p>Intended to be use for streaming out during evaluation of a query.
+     */
+    OutputFormatterCallback<Target> createStreamCallback(
+        PrintStream out, QueryOptions options, @Nullable QueryEnvironment<?> env);
+
+    /**
+     * Same as {@link #createStreamCallback}, but intended to be used for outputting the
+     * already-computed result of a query.
      */
-    OutputFormatterCallback<Target> createStreamCallback(PrintStream out, QueryOptions options);
+    OutputFormatterCallback<Target> createPostFactoStreamCallback(
+        PrintStream out, QueryOptions options);
   }
 
   /**
@@ -209,11 +221,14 @@ public abstract class OutputFormatter implements Serializable {
     }
 
     @Override
-    public void output(QueryOptions options, Digraph<Target> result, PrintStream out,
+    public void output(
+        QueryOptions options,
+        Digraph<Target> result,
+        PrintStream out,
         AspectResolver aspectResolver) throws IOException, InterruptedException {
       setOptions(options, aspectResolver);
       OutputFormatterCallback.processAllTargets(
-          createStreamCallback(out, options), getOrderedTargets(result, options));
+          createPostFactoStreamCallback(out, options), getOrderedTargets(result, options));
     }
   }
 
@@ -235,7 +250,7 @@ public abstract class OutputFormatter implements Serializable {
     }
 
     @Override
-    public OutputFormatterCallback<Target> createStreamCallback(
+    public OutputFormatterCallback<Target> createPostFactoStreamCallback(
         final PrintStream out, final QueryOptions options) {
       return new OutputFormatterCallback<Target>() {
 
@@ -253,6 +268,12 @@ public abstract class OutputFormatter implements Serializable {
         }
       };
     }
+
+    @Override
+    public OutputFormatterCallback<Target> createStreamCallback(
+        PrintStream out, QueryOptions options, QueryEnvironment<?> env) {
+      return createPostFactoStreamCallback(out, options);
+    }
   }
 
   /**
@@ -279,7 +300,7 @@ public abstract class OutputFormatter implements Serializable {
     }
 
     @Override
-    public OutputFormatterCallback<Target> createStreamCallback(
+    public OutputFormatterCallback<Target> createPostFactoStreamCallback(
         final PrintStream out, final QueryOptions options) {
       return new OutputFormatterCallback<Target>() {
         private final Set<String> packageNames = Sets.newTreeSet();
@@ -302,6 +323,12 @@ public abstract class OutputFormatter implements Serializable {
         }
       };
     }
+
+    @Override
+    public OutputFormatterCallback<Target> createStreamCallback(
+        PrintStream out, QueryOptions options, QueryEnvironment<?> env) {
+      return createPostFactoStreamCallback(out, options);
+    }
   }
 
   /**
@@ -318,7 +345,7 @@ public abstract class OutputFormatter implements Serializable {
     }
 
     @Override
-    public OutputFormatterCallback<Target> createStreamCallback(
+    public OutputFormatterCallback<Target> createPostFactoStreamCallback(
         final PrintStream out, final QueryOptions options) {
       return new OutputFormatterCallback<Target>() {
 
@@ -339,6 +366,12 @@ public abstract class OutputFormatter implements Serializable {
         }
       };
     }
+
+    @Override
+    public OutputFormatterCallback<Target> createStreamCallback(
+        PrintStream out, QueryOptions options, QueryEnvironment<?> env) {
+      return createPostFactoStreamCallback(out, options);
+    }
   }
 
   /**
@@ -354,7 +387,7 @@ public abstract class OutputFormatter implements Serializable {
     }
 
     @Override
-    public OutputFormatterCallback<Target> createStreamCallback(
+    public OutputFormatterCallback<Target> createPostFactoStreamCallback(
         final PrintStream out, final QueryOptions options) {
       return new OutputFormatterCallback<Target>() {
         private final Set<Label> printed = CompactHashSet.create();
@@ -412,6 +445,12 @@ public abstract class OutputFormatter implements Serializable {
         }
       };
     }
+
+    @Override
+    public OutputFormatterCallback<Target> createStreamCallback(
+        PrintStream out, QueryOptions options, QueryEnvironment<?> env) {
+      return createPostFactoStreamCallback(out, options);
+    }
   }
 
   private static class RankAndLabel implements Comparable<RankAndLabel> {
diff --git a/src/main/java/com/google/devtools/build/lib/query2/output/ProtoOutputFormatter.java b/src/main/java/com/google/devtools/build/lib/query2/output/ProtoOutputFormatter.java
index c10b6f10be..2cc49f0c5b 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/output/ProtoOutputFormatter.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/output/ProtoOutputFormatter.java
@@ -38,6 +38,7 @@ import com.google.devtools.build.lib.packages.Rule;
 import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.query2.FakeSubincludeTarget;
 import com.google.devtools.build.lib.query2.engine.OutputFormatterCallback;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment;
 import com.google.devtools.build.lib.query2.output.AspectResolver.BuildFileDependencyMode;
 import com.google.devtools.build.lib.query2.output.OutputFormatter.AbstractUnorderedFormatter;
 import com.google.devtools.build.lib.query2.output.QueryOptions.OrderOutput;
@@ -87,7 +88,7 @@ public class ProtoOutputFormatter extends AbstractUnorderedFormatter {
   }
 
   @Override
-  public OutputFormatterCallback<Target> createStreamCallback(
+  public OutputFormatterCallback<Target> createPostFactoStreamCallback(
       final PrintStream out, final QueryOptions options) {
     return new OutputFormatterCallback<Target>() {
 
@@ -114,6 +115,12 @@ public class ProtoOutputFormatter extends AbstractUnorderedFormatter {
     };
   }
 
+  @Override
+  public OutputFormatterCallback<Target> createStreamCallback(
+      PrintStream out, QueryOptions options, QueryEnvironment<?> env) {
+    return createPostFactoStreamCallback(out, options);
+  }
+
   private static Iterable<Target> getSortedLabels(Digraph<Target> result) {
     return Iterables.transform(
         result.getTopologicalOrder(new TargetOrdering()), EXTRACT_NODE_LABEL);
diff --git a/src/main/java/com/google/devtools/build/lib/query2/output/QueryOutputUtils.java b/src/main/java/com/google/devtools/build/lib/query2/output/QueryOutputUtils.java
index 6c63c6bf5a..9d4696b938 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/output/QueryOutputUtils.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/output/QueryOutputUtils.java
@@ -51,7 +51,8 @@ public class QueryOutputUtils {
       StreamedFormatter streamedFormatter = (StreamedFormatter) formatter;
       streamedFormatter.setOptions(queryOptions, aspectResolver);
       OutputFormatterCallback.processAllTargets(
-          streamedFormatter.createStreamCallback(outputStream, queryOptions), targetsResult);
+          streamedFormatter.createPostFactoStreamCallback(outputStream, queryOptions),
+          targetsResult);
     }
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/query2/output/XmlOutputFormatter.java b/src/main/java/com/google/devtools/build/lib/query2/output/XmlOutputFormatter.java
index 9a7aa0cfe0..b60250ebfe 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/output/XmlOutputFormatter.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/output/XmlOutputFormatter.java
@@ -28,6 +28,7 @@ import com.google.devtools.build.lib.packages.Rule;
 import com.google.devtools.build.lib.packages.Target;
 import com.google.devtools.build.lib.query2.FakeSubincludeTarget;
 import com.google.devtools.build.lib.query2.engine.OutputFormatterCallback;
+import com.google.devtools.build.lib.query2.engine.QueryEnvironment;
 import com.google.devtools.build.lib.query2.output.AspectResolver.BuildFileDependencyMode;
 import com.google.devtools.build.lib.query2.output.OutputFormatter.AbstractUnorderedFormatter;
 import com.google.devtools.build.lib.util.Pair;
@@ -61,6 +62,12 @@ class XmlOutputFormatter extends AbstractUnorderedFormatter {
 
   @Override
   public OutputFormatterCallback<Target> createStreamCallback(
+      PrintStream out, QueryOptions options, QueryEnvironment<?> env) {
+    return createPostFactoStreamCallback(out, options);
+  }
+
+  @Override
+  public OutputFormatterCallback<Target> createPostFactoStreamCallback(
       final PrintStream out, final QueryOptions options) {
     return new OutputFormatterCallback<Target>() {
 
diff --git a/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java b/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java
index 565f7a82cf..34afe6be5f 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/commands/QueryCommand.java
@@ -161,7 +161,7 @@ public final class QueryCommand implements BlazeCommand {
         streamedFormatter.setOptions(
             queryOptions,
             queryOptions.aspectDeps.createResolver(env.getPackageManager(), env.getReporter()));
-        callback = streamedFormatter.createStreamCallback(output, queryOptions);
+        callback = streamedFormatter.createStreamCallback(output, queryOptions, queryEnv);
       } else {
         callback = new AggregateAllOutputFormatterCallback<>();
       }