diff --git a/core/trino-main/src/main/java/io/trino/server/CoordinatorModule.java b/core/trino-main/src/main/java/io/trino/server/CoordinatorModule.java index 4629a15c4fd6..adee6e2e7f2b 100644 --- a/core/trino-main/src/main/java/io/trino/server/CoordinatorModule.java +++ b/core/trino-main/src/main/java/io/trino/server/CoordinatorModule.java @@ -200,7 +200,6 @@ protected void setup(Binder binder) // query manager jaxrsBinder(binder).bind(QueryResource.class); - jaxrsBinder(binder).bind(QueryStateInfoResource.class); jaxrsBinder(binder).bind(ResourceGroupStateInfoResource.class); binder.bind(QueryIdGenerator.class).in(Scopes.SINGLETON); binder.bind(SqlQueryManager.class).in(Scopes.SINGLETON); diff --git a/core/trino-main/src/main/java/io/trino/server/QueryStateInfoResource.java b/core/trino-main/src/main/java/io/trino/server/QueryStateInfoResource.java deleted file mode 100644 index dfc00faba788..000000000000 --- a/core/trino-main/src/main/java/io/trino/server/QueryStateInfoResource.java +++ /dev/null @@ -1,122 +0,0 @@ -/* - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package io.trino.server; - -import com.google.inject.Inject; -import io.trino.dispatcher.DispatchManager; -import io.trino.execution.resourcegroups.ResourceGroupInfoProvider; -import io.trino.security.AccessControl; -import io.trino.server.security.ResourceSecurity; -import io.trino.spi.QueryId; -import io.trino.spi.resourcegroups.ResourceGroupId; -import io.trino.spi.security.AccessDeniedException; -import jakarta.servlet.http.HttpServletRequest; -import jakarta.ws.rs.ForbiddenException; -import jakarta.ws.rs.GET; -import jakarta.ws.rs.NotFoundException; -import jakarta.ws.rs.Path; -import jakarta.ws.rs.PathParam; -import jakarta.ws.rs.Produces; -import jakarta.ws.rs.QueryParam; -import jakarta.ws.rs.core.Context; -import jakarta.ws.rs.core.HttpHeaders; -import jakarta.ws.rs.core.MediaType; - -import java.util.List; -import java.util.NoSuchElementException; -import java.util.Optional; -import java.util.regex.Pattern; - -import static com.google.common.base.Strings.isNullOrEmpty; -import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.trino.execution.QueryState.QUEUED; -import static io.trino.security.AccessControlUtil.checkCanViewQueryOwnedBy; -import static io.trino.security.AccessControlUtil.filterQueries; -import static io.trino.server.QueryStateInfo.createQueryStateInfo; -import static io.trino.server.QueryStateInfo.createQueuedQueryStateInfo; -import static io.trino.server.security.ResourceSecurity.AccessType.AUTHENTICATED_USER; -import static java.util.Objects.requireNonNull; - -@Path("/v1/queryState") -@ResourceSecurity(AUTHENTICATED_USER) -public class QueryStateInfoResource -{ - private final DispatchManager dispatchManager; - private final ResourceGroupInfoProvider resourceGroupInfoProvider; - private final AccessControl accessControl; - private final HttpRequestSessionContextFactory sessionContextFactory; - - @Inject - public QueryStateInfoResource( - DispatchManager dispatchManager, - ResourceGroupInfoProvider resourceGroupInfoProvider, - AccessControl accessControl, - HttpRequestSessionContextFactory sessionContextFactory) - { - this.dispatchManager = requireNonNull(dispatchManager, "dispatchManager is null"); - this.resourceGroupInfoProvider = requireNonNull(resourceGroupInfoProvider, "resourceGroupInfoProvider is null"); - this.accessControl = requireNonNull(accessControl, "accessControl is null"); - this.sessionContextFactory = requireNonNull(sessionContextFactory, "sessionContextFactory is null"); - } - - @GET - @Produces(MediaType.APPLICATION_JSON) - public List getQueryStateInfos(@QueryParam("user") String user, @Context HttpServletRequest servletRequest, @Context HttpHeaders httpHeaders) - { - List queryInfos = dispatchManager.getQueries(); - queryInfos = filterQueries(sessionContextFactory.extractAuthorizedIdentity(servletRequest, httpHeaders), queryInfos, accessControl); - - if (!isNullOrEmpty(user)) { - queryInfos = queryInfos.stream() - .filter(queryInfo -> Pattern.matches(user, queryInfo.getSession().getUser())) - .collect(toImmutableList()); - } - - return queryInfos.stream() - .filter(queryInfo -> !queryInfo.getState().isDone()) - .map(this::getQueryStateInfo) - .collect(toImmutableList()); - } - - private QueryStateInfo getQueryStateInfo(BasicQueryInfo queryInfo) - { - Optional groupId = queryInfo.getResourceGroupId(); - if (queryInfo.getState() == QUEUED) { - return createQueuedQueryStateInfo( - queryInfo, - groupId, - groupId.map(group -> resourceGroupInfoProvider.tryGetPathToRoot(group) - .orElseThrow(() -> new IllegalStateException("Resource group not found: " + group)))); - } - return createQueryStateInfo(queryInfo, groupId); - } - - @GET - @Path("{queryId}") - @Produces(MediaType.APPLICATION_JSON) - public QueryStateInfo getQueryStateInfo(@PathParam("queryId") String queryId, @Context HttpServletRequest servletRequest, @Context HttpHeaders httpHeaders) - { - try { - BasicQueryInfo queryInfo = dispatchManager.getQueryInfo(new QueryId(queryId)); - checkCanViewQueryOwnedBy(sessionContextFactory.extractAuthorizedIdentity(servletRequest, httpHeaders), queryInfo.getSession().toIdentity(), accessControl); - return getQueryStateInfo(queryInfo); - } - catch (AccessDeniedException e) { - throw new ForbiddenException(); - } - catch (NoSuchElementException e) { - throw new NotFoundException(); - } - } -} diff --git a/core/trino-main/src/main/java/io/trino/server/ResourceGroupInfo.java b/core/trino-main/src/main/java/io/trino/server/ResourceGroupInfo.java index c8a94f9c4ba7..86345cc2e6e6 100644 --- a/core/trino-main/src/main/java/io/trino/server/ResourceGroupInfo.java +++ b/core/trino-main/src/main/java/io/trino/server/ResourceGroupInfo.java @@ -26,7 +26,7 @@ import static java.util.Objects.requireNonNull; /* - * This class is exposed to external systems via ResourceGroupStateInfoResource and QueryStateInfoResource. + * This class is exposed to external systems via ResourceGroupStateInfoResource. * Be careful while changing it. */ public record ResourceGroupInfo( diff --git a/core/trino-main/src/test/java/io/trino/server/TestQueryStateInfoResource.java b/core/trino-main/src/test/java/io/trino/server/TestQueryStateInfoResource.java deleted file mode 100644 index fe500b439ede..000000000000 --- a/core/trino-main/src/test/java/io/trino/server/TestQueryStateInfoResource.java +++ /dev/null @@ -1,252 +0,0 @@ -/* - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package io.trino.server; - -import com.google.common.io.Closer; -import io.airlift.http.client.HttpClient; -import io.airlift.http.client.Request; -import io.airlift.http.client.UnexpectedResponseException; -import io.airlift.http.client.jetty.JettyHttpClient; -import io.airlift.json.JsonCodec; -import io.airlift.json.JsonCodecFactory; -import io.airlift.json.ObjectMapperProvider; -import io.airlift.units.Duration; -import io.trino.client.QueryResults; -import io.trino.plugin.tpch.TpchPlugin; -import io.trino.server.protocol.spooling.ServerQueryDataJacksonModule; -import io.trino.server.testing.TestingTrinoServer; -import io.trino.spi.ErrorCode; -import org.junit.jupiter.api.AfterAll; -import org.junit.jupiter.api.BeforeAll; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.TestInstance; -import org.junit.jupiter.api.parallel.Execution; - -import java.io.IOException; -import java.util.List; -import java.util.Set; - -import static com.google.common.collect.ImmutableList.toImmutableList; -import static io.airlift.http.client.HttpUriBuilder.uriBuilderFrom; -import static io.airlift.http.client.JsonResponseHandler.createJsonResponseHandler; -import static io.airlift.http.client.Request.Builder.prepareGet; -import static io.airlift.http.client.Request.Builder.preparePost; -import static io.airlift.http.client.StaticBodyGenerator.createStaticBodyGenerator; -import static io.airlift.json.JsonCodec.jsonCodec; -import static io.airlift.json.JsonCodec.listJsonCodec; -import static io.trino.client.ProtocolHeaders.TRINO_HEADERS; -import static io.trino.execution.QueryState.FAILED; -import static io.trino.execution.QueryState.RUNNING; -import static io.trino.server.TestQueryResource.BASIC_QUERY_INFO_CODEC; -import static io.trino.testing.TestingAccessControlManager.TestingPrivilegeType.VIEW_QUERY; -import static io.trino.testing.TestingAccessControlManager.privilege; -import static java.nio.charset.StandardCharsets.UTF_8; -import static java.util.concurrent.TimeUnit.MINUTES; -import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatThrownBy; -import static org.assertj.core.api.Fail.fail; -import static org.junit.jupiter.api.TestInstance.Lifecycle.PER_CLASS; -import static org.junit.jupiter.api.parallel.ExecutionMode.CONCURRENT; - -@TestInstance(PER_CLASS) -@Execution(CONCURRENT) -public class TestQueryStateInfoResource -{ - private static final String LONG_LASTING_QUERY = "SELECT * FROM tpch.sf1.lineitem"; - private static final JsonCodec QUERY_RESULTS_JSON_CODEC = new JsonCodecFactory(new ObjectMapperProvider() - .withModules(Set.of(new ServerQueryDataJacksonModule()))) - .jsonCodec(QueryResults.class); - - private TestingTrinoServer server; - private HttpClient client; - private QueryResults queryResults; - - @BeforeAll - public void setUp() - { - server = TestingTrinoServer.create(); - server.installPlugin(new TpchPlugin()); - server.createCatalog("tpch", "tpch"); - client = new JettyHttpClient(); - - Request request1 = preparePost() - .setUri(uriBuilderFrom(server.getBaseUrl()).replacePath("/v1/statement").build()) - .setBodyGenerator(createStaticBodyGenerator(LONG_LASTING_QUERY, UTF_8)) - .setHeader(TRINO_HEADERS.requestUser(), "user1") - .build(); - queryResults = client.execute(request1, createJsonResponseHandler(QUERY_RESULTS_JSON_CODEC)); - client.execute(prepareGet().setUri(queryResults.getNextUri()).build(), createJsonResponseHandler(QUERY_RESULTS_JSON_CODEC)); - - Request request2 = preparePost() - .setUri(uriBuilderFrom(server.getBaseUrl()).replacePath("/v1/statement").build()) - .setBodyGenerator(createStaticBodyGenerator(LONG_LASTING_QUERY, UTF_8)) - .setHeader(TRINO_HEADERS.requestUser(), "user2") - .build(); - QueryResults queryResults2 = client.execute(request2, createJsonResponseHandler(QUERY_RESULTS_JSON_CODEC)); - client.execute(prepareGet().setUri(queryResults2.getNextUri()).build(), createJsonResponseHandler(QUERY_RESULTS_JSON_CODEC)); - - // queries are started in the background, so they may not all be immediately visible - long start = System.nanoTime(); - while (Duration.nanosSince(start).compareTo(new Duration(5, MINUTES)) < 0) { - List queryInfos = client.execute( - prepareGet() - .setUri(uriBuilderFrom(server.getBaseUrl()).replacePath("/v1/query").build()) - .setHeader(TRINO_HEADERS.requestUser(), "unknown") - .build(), - createJsonResponseHandler(BASIC_QUERY_INFO_CODEC)); - if (queryInfos.size() == 2) { - if (queryInfos.stream().allMatch(info -> info.getState() == RUNNING)) { - break; - } - - List errorCodes = queryInfos.stream() - .filter(info -> info.getState() == FAILED) - .map(BasicQueryInfo::getErrorCode) - .collect(toImmutableList()); - if (!errorCodes.isEmpty()) { - fail("setup queries failed with: " + errorCodes); - } - } - } - } - - @AfterAll - public void tearDown() - throws IOException - { - Closer closer = Closer.create(); - closer.register(server); - closer.register(client); - closer.close(); - server = null; - client = null; - } - - @Test - public void testGetAllQueryStateInfos() - { - List infos = client.execute( - prepareGet() - .setUri(server.resolve("/v1/queryState")) - .setHeader(TRINO_HEADERS.requestUser(), "unknown") - .build(), - createJsonResponseHandler(listJsonCodec(QueryStateInfo.class))); - - assertThat(infos).hasSize(2); - } - - @Test - public void testGetQueryStateInfosForUser() - { - List infos = client.execute( - prepareGet() - .setUri(server.resolve("/v1/queryState?user=user2")) - .setHeader(TRINO_HEADERS.requestUser(), "unknown") - .build(), - createJsonResponseHandler(listJsonCodec(QueryStateInfo.class))); - - assertThat(infos).hasSize(1); - } - - @Test - public void testGetQueryStateInfosForUserNoResult() - { - List infos = client.execute( - prepareGet() - .setUri(server.resolve("/v1/queryState?user=user3")) - .setHeader(TRINO_HEADERS.requestUser(), "unknown") - .build(), - createJsonResponseHandler(listJsonCodec(QueryStateInfo.class))); - - assertThat(infos).isEmpty(); - } - - @Test - public void testGetQueryStateInfo() - { - QueryStateInfo info = client.execute( - prepareGet() - .setUri(server.resolve("/v1/queryState/" + queryResults.getId())) - .setHeader(TRINO_HEADERS.requestUser(), "unknown") - .build(), - createJsonResponseHandler(jsonCodec(QueryStateInfo.class))); - - assertThat(info).isNotNull(); - } - - @Test - public void testGetAllQueryStateInfosDenied() - { - List infos = client.execute( - prepareGet() - .setUri(server.resolve("/v1/queryState")) - .setHeader(TRINO_HEADERS.requestUser(), "any-other-user") - .build(), - createJsonResponseHandler(listJsonCodec(QueryStateInfo.class))); - assertThat(infos).hasSize(2); - - testGetAllQueryStateInfosDenied("user1", 1); - testGetAllQueryStateInfosDenied("any-other-user", 0); - } - - private void testGetAllQueryStateInfosDenied(String executionUser, int expectedCount) - { - server.getAccessControl().deny(privilege(executionUser, "query", VIEW_QUERY)); - try { - List infos = client.execute( - prepareGet() - .setUri(server.resolve("/v1/queryState")) - .setHeader(TRINO_HEADERS.requestUser(), executionUser) - .build(), - createJsonResponseHandler(listJsonCodec(QueryStateInfo.class))); - - assertThat(infos).hasSize(expectedCount); - } - finally { - server.getAccessControl().reset(); - } - } - - @Test - public void testGetQueryStateInfoDenied() - { - server.getAccessControl().deny(privilege("query", VIEW_QUERY)); - try { - assertThatThrownBy(() -> client.execute( - prepareGet() - .setUri(server.resolve("/v1/queryState/" + queryResults.getId())) - .setHeader(TRINO_HEADERS.requestUser(), "unknown") - .build(), - createJsonResponseHandler(jsonCodec(QueryStateInfo.class)))) - .isInstanceOf(UnexpectedResponseException.class) - .matches(throwable -> ((UnexpectedResponseException) throwable).getStatusCode() == 403); - } - finally { - server.getAccessControl().reset(); - } - } - - @Test - public void testGetQueryStateInfoNo() - { - assertThatThrownBy(() -> client.execute( - prepareGet() - .setUri(server.resolve("/v1/queryState/123")) - .setHeader(TRINO_HEADERS.requestUser(), "unknown") - .build(), - createJsonResponseHandler(jsonCodec(QueryStateInfo.class)))) - .isInstanceOf(UnexpectedResponseException.class) - .hasMessageMatching("Expected response code .*, but was 404"); - } -}