Add an afterCommand hook to AbstractBlazeQueryEnvironment.

--
MOS_MIGRATED_REVID=119868282

2 files changed, 5 insertions, 0 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/query2/AbstractBlazeQueryEnvironment.java b/src/main/java/com/google/devtools/build/lib/query2/AbstractBlazeQueryEnvironment.java
index 19430a4736..6a09138efb 100644
--- a/src/main/java/com/google/devtools/build/lib/query2/AbstractBlazeQueryEnvironment.java
+++ b/src/main/java/com/google/devtools/build/lib/query2/AbstractBlazeQueryEnvironment.java
@@ -144,6 +144,9 @@ public abstract class AbstractBlazeQueryEnvironment<T> implements QueryEnvironme
     return new QueryEvalResult(!eventHandler.hasErrors(), empty.get());
   }
 
+  public void afterCommand() {
+  }
+
   public QueryExpression transformParsedQuery(QueryExpression queryExpression) {
     return queryExpression;
   }
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 1302112a9d..cf857fd810 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
@@ -196,6 +196,7 @@ public final class QueryCommand implements BlazeCommand {
       if (!catastrophe) {
         if (streamResults) {
           output.flush();
+          queryEnv.afterCommand();
         }
         try {
           callback.close();
@@ -225,6 +226,7 @@ public final class QueryCommand implements BlazeCommand {
         env.getReporter().handle(Event.error("I/O error: " + e.getMessage()));
         return ExitCode.LOCAL_ENVIRONMENTAL_ERROR;
       } finally {
+        queryEnv.afterCommand();
         // Note that PrintStream#checkError first flushes and then returns whether any
         // error was ever encountered.
         if (output.checkError()) {