Convert DispatchManager.getQueries to return a stream

wendigo · wendigo · commit 679f9c51528a · 2025-11-17T21:14:22.000+01:00
This avoid materialization of potentially large lists
diff --git a/core/trino-main/src/main/java/io/trino/connector/system/QuerySystemTable.java b/core/trino-main/src/main/java/io/trino/connector/system/QuerySystemTable.java
@@ -20,7 +20,6 @@
 import io.trino.execution.QueryInfo;
 import io.trino.execution.QueryStats;
 import io.trino.security.AccessControl;
-import io.trino.server.BasicQueryInfo;
 import io.trino.spi.ErrorCode;
 import io.trino.spi.block.Block;
 import io.trino.spi.block.BlockBuilder;
@@ -104,15 +103,11 @@ public ConnectorTableMetadata getTableMetadata()
     public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, ConnectorSession session, TupleDomain<Integer> constraint)
     {
         checkState(dispatchManager.isPresent(), "Query system table can return results only on coordinator");
-
-        List<BasicQueryInfo> queries = dispatchManager.get().getQueries();
-        queries = filterQueries(((FullConnectorSession) session).getSession().getIdentity(), queries, accessControl);
-
         Builder table = InMemoryRecordSet.builder(QUERY_TABLE);
-        for (BasicQueryInfo queryInfo : queries) {
+        filterQueries(dispatchManager.get()::getQueries, ((FullConnectorSession) session).getSession().getIdentity(), accessControl).forEach(queryInfo -> {
             Optional<QueryInfo> fullQueryInfo = dispatchManager.get().getFullQueryInfo(queryInfo.getQueryId());
             if (fullQueryInfo.isEmpty()) {
-                continue;
+                return;
             }
             QueryStats queryStats = fullQueryInfo.get().getQueryStats();
             table.addRow(
@@ -134,7 +129,7 @@ public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, Connect
 
                     Optional.ofNullable(queryInfo.getErrorType()).map(Enum::name).orElse(null),
                     Optional.ofNullable(queryInfo.getErrorCode()).map(ErrorCode::getName).orElse(null));
-        }
+        });
         return table.build().cursor();
     }
 
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
@@ -51,13 +51,12 @@
 import org.weakref.jmx.Managed;
 import org.weakref.jmx.Nested;
 
-import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.stream.Stream;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.collect.ImmutableList.toImmutableList;
 import static io.airlift.concurrent.Threads.daemonThreadsNamed;
 import static io.trino.execution.QueryState.FINISHING;
 import static io.trino.execution.QueryState.QUEUED;
@@ -311,11 +310,10 @@ public ListenableFuture<Void> waitForDispatched(QueryId queryId)
                 .orElseGet(Futures::immediateVoidFuture);
     }
 
-    public List<BasicQueryInfo> getQueries()
+    public Stream<BasicQueryInfo> getQueries()
     {
         return queryTracker.getAllQueries()
-                .map(DispatchQuery::getBasicQueryInfo)
-                .collect(toImmutableList());
+                .map(DispatchQuery::getBasicQueryInfo);
     }
 
     @Managed
diff --git a/core/trino-main/src/main/java/io/trino/security/AccessControlUtil.java b/core/trino-main/src/main/java/io/trino/security/AccessControlUtil.java
@@ -13,12 +13,13 @@
  */
 package io.trino.security;
 
+import com.google.common.base.Supplier;
 import io.trino.SessionRepresentation;
 import io.trino.server.BasicQueryInfo;
 import io.trino.spi.security.Identity;
 
 import java.util.Collection;
-import java.util.List;
+import java.util.stream.Stream;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 
@@ -34,22 +35,23 @@ public static void checkCanViewQueryOwnedBy(Identity identity, Identity queryOwn
         accessControl.checkCanViewQueryOwnedBy(identity, queryOwner);
     }
 
-    public static List<BasicQueryInfo> filterQueries(Identity identity, List<BasicQueryInfo> queries, AccessControl accessControl)
+    public static Stream<BasicQueryInfo> filterQueries(Supplier<Stream<BasicQueryInfo>> queries, Identity identity, AccessControl accessControl)
     {
-        Collection<Identity> owners = queries.stream()
+        Collection<Identity> owners = queries
+                .get()
                 .map(BasicQueryInfo::getSession)
                 .map(SessionRepresentation::toIdentity)
                 .filter(owner -> !owner.getUser().equals(identity.getUser()))
                 .distinct()
                 .collect(toImmutableList());
         Collection<Identity> allowedOwners = accessControl.filterQueriesOwnedBy(identity, owners);
 
-        return queries.stream()
+        return queries
+                .get()
                 .filter(queryInfo -> {
                     Identity queryIdentity = queryInfo.getSession().toIdentity();
                     return queryIdentity.getUser().equals(identity.getUser()) || allowedOwners.contains(queryIdentity);
-                })
-                .collect(toImmutableList());
+                });
     }
 
     public static void checkCanKillQueryOwnedBy(Identity identity, Identity queryOwner, AccessControl accessControl)
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
@@ -36,13 +36,12 @@
 import jakarta.ws.rs.core.Response;
 import jakarta.ws.rs.core.Response.Status;
 
-import java.util.List;
 import java.util.Locale;
 import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.Set;
+import java.util.stream.Stream;
 
-import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static io.trino.connector.system.KillQueryProcedure.createKillQueryException;
 import static io.trino.connector.system.KillQueryProcedure.createPreemptQueryException;
@@ -73,22 +72,18 @@ public QueryResource(DispatchManager dispatchManager, AccessControl accessContro
     }
 
     @GET
-    public List<BasicQueryInfo> getAllQueryInfo(@QueryParam("state") Set<String> stateFilters, @Context HttpServletRequest servletRequest, @Context HttpHeaders httpHeaders)
+    public Stream<BasicQueryInfo> getAllQueryInfo(@QueryParam("state") Set<String> stateFilters, @Context HttpServletRequest servletRequest, @Context HttpHeaders httpHeaders)
     {
         Set<QueryState> expectedStates = stateFilters.stream()
                 .map(state -> state.toUpperCase(Locale.ENGLISH))
                 .map(QueryState::valueOf)
                 .collect(toImmutableSet());
 
-        List<BasicQueryInfo> queries = dispatchManager.getQueries();
-        queries = filterQueries(sessionContextFactory.extractAuthorizedIdentity(servletRequest, httpHeaders), queries, accessControl);
-
+        Stream<BasicQueryInfo> queries = filterQueries(dispatchManager::getQueries, sessionContextFactory.extractAuthorizedIdentity(servletRequest, httpHeaders), accessControl);
         if (expectedStates.isEmpty()) {
             return queries;
         }
-        return queries.stream()
-                .filter(query -> expectedStates.contains(query.getState()))
-                .collect(toImmutableList());
+        return queries.filter(query -> expectedStates.contains(query.getState()));
     }
 
     @GET
diff --git a/core/trino-main/src/main/java/io/trino/server/ui/ClusterStatsResource.java b/core/trino-main/src/main/java/io/trino/server/ui/ClusterStatsResource.java
@@ -30,6 +30,8 @@
 import jakarta.ws.rs.Produces;
 import jakarta.ws.rs.core.MediaType;
 
+import java.util.Iterator;
+
 import static com.google.common.math.DoubleMath.roundToLong;
 import static io.trino.server.security.ResourceSecurity.AccessType.WEB_UI;
 import static java.math.RoundingMode.HALF_UP;
@@ -77,7 +79,9 @@ public ClusterStats getClusterStats()
         long totalInputBytes = dispatchManager.getStats().getConsumedInputBytes().getTotalCount();
         long totalCpuTimeSecs = dispatchManager.getStats().getConsumedCpuTimeSecs().getTotalCount();
 
-        for (BasicQueryInfo query : dispatchManager.getQueries()) {
+        Iterator<BasicQueryInfo> iterator = dispatchManager.getQueries().iterator();
+        while (iterator.hasNext()) {
+            BasicQueryInfo query = iterator.next();
             if (query.getState() == QueryState.QUEUED) {
                 queuedQueries++;
             }
diff --git a/core/trino-main/src/main/java/io/trino/server/ui/UiQueryResource.java b/core/trino-main/src/main/java/io/trino/server/ui/UiQueryResource.java
@@ -16,7 +16,6 @@
 import com.fasterxml.jackson.annotation.JsonTypeInfo;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.cfg.ContextAttributes;
-import com.google.common.collect.ImmutableList;
 import com.google.inject.Inject;
 import io.airlift.json.JsonCodec;
 import io.airlift.json.JsonCodecFactory;
@@ -47,10 +46,10 @@
 import jakarta.ws.rs.core.Response;
 import jakarta.ws.rs.core.Response.Status;
 
-import java.util.List;
 import java.util.Locale;
 import java.util.NoSuchElementException;
 import java.util.Optional;
+import java.util.stream.Stream;
 
 import static com.fasterxml.jackson.annotation.JsonIgnoreProperties.Value.forIgnoredProperties;
 import static io.trino.connector.system.KillQueryProcedure.createKillQueryException;
@@ -86,20 +85,16 @@ public UiQueryResource(ObjectMapper objectMapper, DispatchManager dispatchManage
     }
 
     @GET
-    public List<TrimmedBasicQueryInfo> getAllQueryInfo(@QueryParam("state") String stateFilter, @Context HttpServletRequest servletRequest, @Context HttpHeaders httpHeaders)
+    public Stream<TrimmedBasicQueryInfo> getAllQueryInfo(@QueryParam("state") String stateFilter, @Context HttpServletRequest servletRequest, @Context HttpHeaders httpHeaders)
     {
         QueryState expectedState = stateFilter == null ? null : QueryState.valueOf(stateFilter.toUpperCase(Locale.ENGLISH));
-
-        List<BasicQueryInfo> queries = dispatchManager.getQueries();
-        queries = filterQueries(sessionContextFactory.extractAuthorizedIdentity(servletRequest, httpHeaders), queries, accessControl);
-
-        ImmutableList.Builder<TrimmedBasicQueryInfo> builder = ImmutableList.builder();
-        for (BasicQueryInfo queryInfo : queries) {
+        Stream<BasicQueryInfo> queries = filterQueries(dispatchManager::getQueries, sessionContextFactory.extractAuthorizedIdentity(servletRequest, httpHeaders), accessControl);
+        return queries.flatMap(queryInfo -> {
             if (stateFilter == null || queryInfo.getState() == expectedState) {
-                builder.add(new TrimmedBasicQueryInfo(queryInfo));
+                return Stream.of(new TrimmedBasicQueryInfo(queryInfo));
             }
-        }
-        return builder.build();
+            return Stream.empty();
+        });
     }
 
     @GET
diff --git a/testing/trino-tests/src/test/java/io/trino/execution/QueryRunnerUtil.java b/testing/trino-tests/src/test/java/io/trino/execution/QueryRunnerUtil.java
@@ -17,7 +17,6 @@
 import io.opentelemetry.api.trace.Span;
 import io.trino.Session;
 import io.trino.dispatcher.DispatchManager;
-import io.trino.server.BasicQueryInfo;
 import io.trino.server.SessionContext;
 import io.trino.server.protocol.Slug;
 import io.trino.spi.QueryId;
@@ -55,13 +54,14 @@ public static void waitForQueryState(QueryRunner queryRunner, QueryId queryId, S
             throws InterruptedException
     {
         DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager();
+        QueryManager queryManager = queryRunner.getCoordinator().getQueryManager();
         do {
             // Heartbeat all the running queries, so they don't die while we're waiting
-            for (BasicQueryInfo queryInfo : dispatchManager.getQueries()) {
+            dispatchManager.getQueries().forEach(queryInfo -> {
                 if (queryInfo.getState() == RUNNING) {
-                    dispatchManager.getQueryInfo(queryInfo.getQueryId());
+                    queryManager.recordHeartbeat(queryInfo.getQueryId());
                 }
-            }
+            });
             MILLISECONDS.sleep(100);
         }
         while (!expectedQueryStates.contains(dispatchManager.getQueryInfo(queryId).getState()));
diff --git a/testing/trino-tests/src/test/java/io/trino/security/TestSessionImpersonation.java b/testing/trino-tests/src/test/java/io/trino/security/TestSessionImpersonation.java
@@ -139,12 +139,11 @@ private void assertSessionUsersAndGroups(
 
     private BasicQueryInfo getQueryInfo(String query)
     {
-        QueryId queryId = null;
-        for (BasicQueryInfo basicQueryInfo : server.getDispatchManager().getQueries()) {
-            if (basicQueryInfo.getQuery().equals(query)) {
-                queryId = basicQueryInfo.getQueryId();
-            }
-        }
+        QueryId queryId = server.getDispatchManager().getQueries()
+                .filter(basicQueryInfo -> basicQueryInfo.getQuery().equals(query))
+                .map(BasicQueryInfo::getQueryId)
+                .findFirst()
+                .orElseThrow(() -> new IllegalStateException("Could not find query"));
         return server.getDispatchManager().getQueryInfo(queryId);
     }
 
