Add bazel query --output_file option, which writes query results directly to a file

This is a proposed fix for #24293

This speeds up a fully warm `bazel query ...` by 23.7%, reducing wall time from 1m49s to 1m23s

```
$ time bazel query '...' --output=streamed_proto > queryoutput4.streamedproto

real    1m48.768s
user    0m27.410s
sys     0m19.646s

$ time bazel query '...' --output=streamed_proto --output_file=queryoutput5.streamedproto

real    1m22.920s
user    0m0.045s
sys     0m0.016s
```

_💁‍♂️ Note: when combined with #24305, total wall time is 37s, an overall reduction of 66%._

Closes #24298.

PiperOrigin-RevId: 700583890
Change-Id: Ic13f0611aca60c2ce8641e72a0fcfc330f13c803

Author: keithl-stripe

scripts/bash_completion_test.sh

@@ -452,7 +452,7 @@ test_build_options() {
 
 test_query_options() {
     assert_expansion 'query --out' \
-                     'query --output='
+                     'query --output'
 
     # Basic label expansion works for query, too.
     make_packages

src/main/java/com/google/devtools/build/lib/BUILD

@@ -403,6 +403,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/query2",
         "//src/main/java/com/google/devtools/build/lib/query2:aquery-utils",
         "//src/main/java/com/google/devtools/build/lib/query2/common:cquery-node",
+        "//src/main/java/com/google/devtools/build/lib/query2/common:options",
         "//src/main/java/com/google/devtools/build/lib/query2/engine",
         "//src/main/java/com/google/devtools/build/lib/query2/query/output",
         "//src/main/java/com/google/devtools/build/lib/rules:repository/repository_function",

src/main/java/com/google/devtools/build/lib/buildtool/AqueryProcessor.java

@@ -67,11 +67,16 @@ public AqueryProcessor(
     actionFilters = buildActionFilters(queryExpression);
   }
 
+  @Override
+  protected AqueryOptions getQueryOptions(CommandEnvironment env) {
+    return env.getOptions().getOptions(AqueryOptions.class);
+  }
+
   /** Outputs the current action graph from Skyframe. */
   public BlazeCommandResult dumpActionGraphFromSkyframe(CommandEnvironment env) {
+    AqueryOptions aqueryOptions = getQueryOptions(env);
     try (QueryRuntimeHelper queryRuntimeHelper =
-        env.getRuntime().getQueryRuntimeHelperFactory().create(env)) {
-      AqueryOptions aqueryOptions = env.getOptions().getOptions(AqueryOptions.class);
+        env.getRuntime().getQueryRuntimeHelperFactory().create(env, aqueryOptions)) {
 
       PrintStream printStream =
           queryRuntimeHelper.getOutputStreamForQueryOutput() == null

src/main/java/com/google/devtools/build/lib/buildtool/CqueryProcessor.java

@@ -19,6 +19,7 @@
 import com.google.devtools.build.lib.cmdline.TargetPattern;
 import com.google.devtools.build.lib.packages.semantics.BuildLanguageOptions;
 import com.google.devtools.build.lib.query2.PostAnalysisQueryEnvironment.TopLevelConfigurations;
+import com.google.devtools.build.lib.query2.common.CommonQueryOptions;
 import com.google.devtools.build.lib.query2.common.CqueryNode;
 import com.google.devtools.build.lib.query2.cquery.ConfiguredTargetQueryEnvironment;
 import com.google.devtools.build.lib.query2.cquery.CqueryOptions;
@@ -36,6 +37,11 @@ public CqueryProcessor(
     super(queryExpression, mainRepoTargetParser);
   }
 
+  @Override
+  protected CommonQueryOptions getQueryOptions(CommandEnvironment env) {
+    return env.getOptions().getOptions(CqueryOptions.class);
+  }
+
   @Override
   protected ConfiguredTargetQueryEnvironment getQueryEnvironment(
       BuildRequest request,

src/main/java/com/google/devtools/build/lib/buildtool/PostAnalysisQueryProcessor.java

@@ -26,6 +26,7 @@
 import com.google.devtools.build.lib.query2.NamedThreadSafeOutputFormatterCallback;
 import com.google.devtools.build.lib.query2.PostAnalysisQueryEnvironment;
 import com.google.devtools.build.lib.query2.PostAnalysisQueryEnvironment.TopLevelConfigurations;
+import com.google.devtools.build.lib.query2.common.CommonQueryOptions;
 import com.google.devtools.build.lib.query2.engine.QueryEvalResult;
 import com.google.devtools.build.lib.query2.engine.QueryException;
 import com.google.devtools.build.lib.query2.engine.QueryExpression;
@@ -90,7 +91,7 @@ public void process(
       }
 
       try (QueryRuntimeHelper queryRuntimeHelper =
-          env.getRuntime().getQueryRuntimeHelperFactory().create(env)) {
+          env.getRuntime().getQueryRuntimeHelperFactory().create(env, getQueryOptions(env))) {
         doPostAnalysisQuery(
             request,
             env,
@@ -131,6 +132,8 @@ public void process(
     }
   }
 
+  protected abstract CommonQueryOptions getQueryOptions(CommandEnvironment env);
+
   protected abstract PostAnalysisQueryEnvironment<T> getQueryEnvironment(
       BuildRequest request,
       CommandEnvironment env,

src/main/java/com/google/devtools/build/lib/query2/common/CommonQueryOptions.java

@@ -373,6 +373,10 @@ public AspectResolutionModeConverter() {
               + "applicable to --output=graph.")
   public boolean graphFactored;
 
+  ///////////////////////////////////////////////////////////
+  // INPUT / OUTPUT OPTIONS                                //
+  ///////////////////////////////////////////////////////////
+
   @Option(
       name = "query_file",
       defaultValue = "",
@@ -382,4 +386,15 @@ public AspectResolutionModeConverter() {
           "If set, query will read the query from the file named here, rather than on the command "
               + "line. It is an error to specify a file here as well as a command-line query.")
   public String queryFile;
+
+  @Option(
+      name = "output_file",
+      defaultValue = "",
+      documentationCategory = OptionDocumentationCategory.QUERY,
+      effectTags = {OptionEffectTag.TERMINAL_OUTPUT},
+      help =
+          "When specified, query results will be written directly to this file, and nothing will be"
+              + " printed to Bazel's standard output stream (stdout). In benchmarks, this is"
+              + " generally faster than <code>bazel query &gt; file</code>.")
+  public String outputFile;
 }

src/main/java/com/google/devtools/build/lib/runtime/QueryRuntimeHelper.java

@@ -15,10 +15,14 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
+import com.google.common.base.Strings;
+import com.google.devtools.build.lib.query2.common.CommonQueryOptions;
 import com.google.devtools.build.lib.server.FailureDetails;
 import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
 import com.google.devtools.build.lib.server.FailureDetails.Query;
 import com.google.devtools.build.lib.server.FailureDetails.Query.Code;
+import com.google.devtools.build.lib.vfs.Path;
+import java.io.IOException;
 import java.io.OutputStream;
 
 /**
@@ -47,14 +51,17 @@ public interface QueryRuntimeHelper extends AutoCloseable {
 
   /** Factory for {@link QueryRuntimeHelper} instances. */
   interface Factory {
-    QueryRuntimeHelper create(CommandEnvironment env) throws QueryRuntimeHelperException;
+    QueryRuntimeHelper create(CommandEnvironment env, CommonQueryOptions options)
+        throws QueryRuntimeHelperException;
   }
 
   /**
    * A {@link Factory} for {@link StdoutQueryRuntimeHelper} instances that simply wrap the given
    * {@link CommandEnvironment} instance's stdout.
    *
    * <p>This is intended to be the default {@link Factory}.
+   *
+   * <p>If {@code --output_file} is set, the stdout is redirected to the defined path value instead.
    */
   class StdoutQueryRuntimeHelperFactory implements Factory {
     public static final StdoutQueryRuntimeHelperFactory INSTANCE =
@@ -63,8 +70,14 @@ class StdoutQueryRuntimeHelperFactory implements Factory {
     private StdoutQueryRuntimeHelperFactory() {}
 
     @Override
-    public QueryRuntimeHelper create(CommandEnvironment env) {
-      return createInternal(env.getReporter().getOutErr().getOutputStream());
+    public QueryRuntimeHelper create(CommandEnvironment env, CommonQueryOptions options)
+        throws QueryRuntimeHelperException {
+      if (Strings.isNullOrEmpty(options.outputFile)) {
+        return createInternal(env.getReporter().getOutErr().getOutputStream());
+      } else {
+        return FileQueryRuntimeHelper.create(
+            env.getWorkingDirectory().getRelative(options.outputFile));
+      }
     }
 
     public QueryRuntimeHelper createInternal(OutputStream stdoutOutputStream) {
@@ -91,6 +104,49 @@ public void afterQueryOutputIsWritten() {}
       @Override
       public void close() {}
     }
+
+    /**
+     * A {@link QueryRuntimeHelper} that wraps a {@link java.io.FileOutputStream} instead of writing
+     * to standard out, for improved performance.
+     */
+    public static class FileQueryRuntimeHelper implements QueryRuntimeHelper {
+      private final Path path;
+      private final OutputStream out;
+
+      private FileQueryRuntimeHelper(Path path) throws IOException {
+        this.path = path;
+        this.out = path.getOutputStream();
+      }
+
+      public static FileQueryRuntimeHelper create(Path path) throws QueryRuntimeHelperException {
+        try {
+          return new FileQueryRuntimeHelper(path);
+        } catch (IOException e) {
+          throw new QueryRuntimeHelperException(
+              "Could not open query output file " + path.getPathString(),
+              Code.QUERY_OUTPUT_WRITE_FAILURE,
+              e);
+        }
+      }
+
+      @Override
+      public OutputStream getOutputStreamForQueryOutput() {
+        return out;
+      }
+
+      @Override
+      public void afterQueryOutputIsWritten() {}
+
+      @Override
+      public void close() throws QueryRuntimeHelperException {
+        try {
+          out.close();
+        } catch (IOException e) {
+          throw new QueryRuntimeHelperException(
+              "Could not close query output file " + path, Code.QUERY_OUTPUT_WRITE_FAILURE, e);
+        }
+      }
+    }
   }
 
   /** Describes what went wrong in {@link QueryRuntimeHelper}. */

src/main/java/com/google/devtools/build/lib/runtime/commands/QueryEnvironmentBasedCommand.java

@@ -167,7 +167,7 @@ private BlazeCommandResult execInternal(CommandEnvironment env, OptionsParsingRe
             .getLabelPrinter(starlarkSemantics, mainRepoTargetParser.getRepoMapping());
 
     try (QueryRuntimeHelper queryRuntimeHelper =
-        env.getRuntime().getQueryRuntimeHelperFactory().create(env)) {
+        env.getRuntime().getQueryRuntimeHelperFactory().create(env, queryOptions)) {
       Either<BlazeCommandResult, QueryEvalResult> result;
       try (AbstractBlazeQueryEnvironment<Target> queryEnv =
           newQueryEnvironment(

src/test/shell/integration/bazel_query_test.sh

@@ -42,7 +42,7 @@ fi
 source "$(rlocation "io_bazel/src/test/shell/integration_test_setup.sh")" \
   || { echo "integration_test_setup.sh not found!" >&2; exit 1; }
 
-# `uname` returns the current platform, e.g "MSYS_NT-10.0" or "Linux".
+# `uname` returns the current platform, e.g. "MSYS_NT-10.0" or "Linux".
 # `tr` converts all upper case letters to lower case.
 # `case` matches the result if the `uname | tr` expression to string prefixes
 # that use the same wildcards as names do in Bash, i.e. "msys*" matches strings
@@ -82,6 +82,22 @@ EOF
   expect_log "//peach:harken"
 }
 
+function test_output_to_file() {
+  rm -rf peach
+  mkdir -p peach
+  cat > peach/BUILD <<EOF
+sh_library(name='brighton', deps=[':harken'])
+sh_library(name='harken')
+EOF
+
+  bazel query 'deps(//peach:brighton)' --output_file=$TEST_log > $TEST_TMPDIR/query_stdout
+
+  expect_log "//peach:brighton"
+  expect_log "//peach:harken"
+
+  assert_equals "" "$(<$TEST_TMPDIR/query_stdout)"
+}
+
 function test_invalid_query_fails_parsing() {
   bazel query 'deps("--bad_target_name_from_bad_script")' >& "$TEST_log" \
     && fail "Expected failure"