Support multiple state filters on getAllQueryInfo API

trinodb · Dec 12, 2024 · cd15023 · cd15023
1 parent da79aed
commit cd15023
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 10 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/server/QueryResource.java b/core/trino-main/src/main/java/io/trino/server/QueryResource.java
@@ -13,7 +13,6 @@
  */
 package io.trino.server;
 
-import com.google.common.collect.ImmutableList;
 import com.google.inject.Inject;
 import io.trino.dispatcher.DispatchManager;
 import io.trino.execution.QueryInfo;
@@ -41,6 +40,8 @@
 import java.util.Locale;
 import java.util.NoSuchElementException;
 import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
 
 import static io.trino.connector.system.KillQueryProcedure.createKillQueryException;
 import static io.trino.connector.system.KillQueryProcedure.createPreemptQueryException;
@@ -71,20 +72,19 @@ public QueryResource(DispatchManager dispatchManager, AccessControl accessContro
 
     @ResourceSecurity(AUTHENTICATED_USER)
     @GET
-    public List<BasicQueryInfo> getAllQueryInfo(@QueryParam("state") String stateFilter, @Context HttpServletRequest servletRequest, @Context HttpHeaders httpHeaders)
+    public List<BasicQueryInfo> getAllQueryInfo(@QueryParam("state") Set<String> stateFilters, @Context HttpServletRequest servletRequest, @Context HttpHeaders httpHeaders)
     {
-        QueryState expectedState = stateFilter == null ? null : QueryState.valueOf(stateFilter.toUpperCase(Locale.ENGLISH));
+        Set<QueryState> expectedStates = stateFilters.stream()
+                .map(f -> f.toUpperCase(Locale.ENGLISH))
+                .map(QueryState::valueOf)
+                .collect(Collectors.toUnmodifiableSet());
 
         List<BasicQueryInfo> queries = dispatchManager.getQueries();
         queries = filterQueries(sessionContextFactory.extractAuthorizedIdentity(servletRequest, httpHeaders), queries, accessControl);
 
-        ImmutableList.Builder<BasicQueryInfo> builder = ImmutableList.builder();
-        for (BasicQueryInfo queryInfo : queries) {
-            if (stateFilter == null || queryInfo.getState() == expectedState) {
-                builder.add(queryInfo);
-            }
-        }
-        return builder.build();
+        return queries.stream()
+                .filter(q -> expectedStates.isEmpty() || expectedStates.contains(q.getState()))
+                .collect(Collectors.toUnmodifiableList());
     }
 
     @ResourceSecurity(AUTHENTICATED_USER)

diff --git a/core/trino-main/src/test/java/io/trino/server/TestQueryResource.java b/core/trino-main/src/test/java/io/trino/server/TestQueryResource.java
@@ -163,12 +163,17 @@ public void testGetQueryInfos()
         assertThat(infos).isEmpty();
         assertStateCounts(infos, 0, 0, 0);
 
+        infos = getQueryInfos("/v1/query?state=finished&state=failed&state=running");
+        assertThat(infos).hasSize(3);
+        assertStateCounts(infos, 2, 1, 0);
+
         server.getAccessControl().deny(privilege("query", VIEW_QUERY));
         try {
             assertThat(getQueryInfos("/v1/query")).isEmpty();
             assertThat(getQueryInfos("/v1/query?state=finished")).isEmpty();
             assertThat(getQueryInfos("/v1/query?state=failed")).isEmpty();
             assertThat(getQueryInfos("/v1/query?state=running")).isEmpty();
+            assertThat(getQueryInfos("/v1/query?state=finished&state=failed&state=running")).isEmpty();
         }
         finally {
             server.getAccessControl().reset();