Ensure cached query results are discarded

Send next token one more time to make the client acknowledge results to discard cached query results
trinodb · Nov 4, 2021 · da2cd0a · da2cd0a
1 parent 12b3efc
commit da2cd0a
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/server/protocol/Query.java b/core/trino-main/src/main/java/io/trino/server/protocol/Query.java
@@ -417,6 +417,8 @@ private synchronized QueryResults getNextResult(long token, UriInfo uriInfo, Dat
         QueryInfo queryInfo = queryManager.getFullQueryInfo(queryId);
         queryManager.recordHeartbeat(queryId);
 
+        closeExchangeClientIfNecessary(queryInfo);
+
         // fetch result data from exchange
         QueryResultRows resultRows = removePagesFromExchange(queryInfo, targetResultSize.toBytes());
 
@@ -426,14 +428,12 @@ private synchronized QueryResults getNextResult(long token, UriInfo uriInfo, Dat
             updateCount = updatedRowsCount.orElse(null);
         }
 
-        closeExchangeClientIfNecessary(queryInfo);
-
         // advance next token
         // only return a next if
         // (1) the query is not done AND the query state is not FAILED
         //   OR
         // (2)there is more data to send (due to buffering)
-        if ((!queryInfo.isFinalQueryInfo() && queryInfo.getState() != FAILED) || !exchangeClient.isClosed()) {
+        if (queryInfo.getState() != FAILED && (!queryInfo.isFinalQueryInfo() || !exchangeClient.isClosed() || (lastResult != null && lastResult.getData() != null))) {
             nextToken = OptionalLong.of(token + 1);
         }
         else {