diff --git a/core/trino-main/src/main/java/io/trino/dispatcher/DispatchManager.java b/core/trino-main/src/main/java/io/trino/dispatcher/DispatchManager.java index 8374b3007bc2..7e9e620c2fa1 100644 --- a/core/trino-main/src/main/java/io/trino/dispatcher/DispatchManager.java +++ b/core/trino-main/src/main/java/io/trino/dispatcher/DispatchManager.java @@ -146,7 +146,7 @@ public ListenableFuture createQuery(QueryId queryId, Span querySpan, Slug requireNonNull(sessionContext, "sessionContext is null"); requireNonNull(query, "query is null"); checkArgument(!query.isEmpty(), "query must not be empty string"); - checkArgument(queryTracker.tryGetQuery(queryId).isEmpty(), "query %s already exists", queryId); + checkArgument(!queryTracker.hasQuery(queryId), "query %s already exists", queryId); // It is important to return a future implementation which ignores cancellation request. // Using NonCancellationPropagatingFuture is not enough; it does not propagate cancel to wrapped future @@ -295,7 +295,7 @@ public long getRunningQueries() public boolean isQueryRegistered(QueryId queryId) { - return queryTracker.tryGetQuery(queryId).isPresent(); + return queryTracker.hasQuery(queryId); } public DispatchQuery getQuery(QueryId queryId) diff --git a/core/trino-main/src/main/java/io/trino/execution/QueryManager.java b/core/trino-main/src/main/java/io/trino/execution/QueryManager.java index c0d3c588e389..af52ae4edbeb 100644 --- a/core/trino-main/src/main/java/io/trino/execution/QueryManager.java +++ b/core/trino-main/src/main/java/io/trino/execution/QueryManager.java @@ -92,6 +92,8 @@ QueryInfo getFullQueryInfo(QueryId queryId) QueryState getQueryState(QueryId queryId) throws NoSuchElementException; + boolean hasQuery(QueryId queryId); + /** * Updates the client heartbeat time, to prevent the query from be automatically purged. * If the query does not exist, the call is ignored. diff --git a/core/trino-main/src/main/java/io/trino/execution/QueryTracker.java b/core/trino-main/src/main/java/io/trino/execution/QueryTracker.java index 6d4d222f5679..afd81baba987 100644 --- a/core/trino-main/src/main/java/io/trino/execution/QueryTracker.java +++ b/core/trino-main/src/main/java/io/trino/execution/QueryTracker.java @@ -147,6 +147,12 @@ public T getQuery(QueryId queryId) .orElseThrow(() -> new NoSuchElementException(queryId.toString())); } + public boolean hasQuery(QueryId queryId) + { + requireNonNull(queryId, "queryId is null"); + return queries.containsKey(queryId); + } + public Optional tryGetQuery(QueryId queryId) { requireNonNull(queryId, "queryId is null"); diff --git a/core/trino-main/src/main/java/io/trino/execution/SqlQueryManager.java b/core/trino-main/src/main/java/io/trino/execution/SqlQueryManager.java index 0e206a349260..01424b3f95b0 100644 --- a/core/trino-main/src/main/java/io/trino/execution/SqlQueryManager.java +++ b/core/trino-main/src/main/java/io/trino/execution/SqlQueryManager.java @@ -226,6 +226,12 @@ public QueryState getQueryState(QueryId queryId) return queryTracker.getQuery(queryId).getState(); } + @Override + public boolean hasQuery(QueryId queryId) + { + return queryTracker.hasQuery(queryId); + } + @Override public void recordHeartbeat(QueryId queryId) { diff --git a/core/trino-main/src/main/java/io/trino/server/protocol/ExecutingStatementResource.java b/core/trino-main/src/main/java/io/trino/server/protocol/ExecutingStatementResource.java index 307a76d97ce3..d8bd9204b6de 100644 --- a/core/trino-main/src/main/java/io/trino/server/protocol/ExecutingStatementResource.java +++ b/core/trino-main/src/main/java/io/trino/server/protocol/ExecutingStatementResource.java @@ -120,11 +120,7 @@ public ExecutingStatementResource( try { for (QueryId queryId : queries.keySet()) { // forget about this query if the query manager is no longer tracking it - try { - queryManager.getQueryState(queryId); - } - catch (NoSuchElementException e) { - // query is no longer registered + if (!queryManager.hasQuery(queryId)) { Query query = queries.remove(queryId); if (query != null) { query.dispose();