From 1fc2f6e4b08d92f0c68d31afd7890ea9c211d4c2 Mon Sep 17 00:00:00 2001 From: Tejaswini Bandlamudi <96047043+tejaswini-imply@users.noreply.github.com> Date: Fri, 24 Jun 2022 09:21:25 +0530 Subject: [PATCH] Throw BadQueryContextException if context params cannot be parsed (#12680) --- .../druid/query/BadQueryContextException.java | 44 +++++++++++++++++++ .../org/apache/druid/query/QueryContexts.java | 11 +++-- .../apache/druid/query/QueryContextsTest.java | 15 +++++++ .../apache/druid/server/QueryResource.java | 3 +- .../apache/druid/sql/http/SqlResource.java | 3 +- .../druid/sql/http/SqlResourceTest.java | 25 +++++++++++ 6 files changed, 94 insertions(+), 7 deletions(-) create mode 100644 processing/src/main/java/org/apache/druid/query/BadQueryContextException.java diff --git a/processing/src/main/java/org/apache/druid/query/BadQueryContextException.java b/processing/src/main/java/org/apache/druid/query/BadQueryContextException.java new file mode 100644 index 00000000000..1991656332c --- /dev/null +++ b/processing/src/main/java/org/apache/druid/query/BadQueryContextException.java @@ -0,0 +1,44 @@ +/* + * 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 com.fasterxml.jackson.annotation.JsonProperty; + +public class BadQueryContextException extends BadQueryException +{ + public static final String ERROR_CODE = "Query context parse failed"; + public static final String ERROR_CLASS = BadQueryContextException.class.getName(); + + public BadQueryContextException(Exception e) + { + this(ERROR_CODE, e.getMessage(), ERROR_CLASS); + } + + @JsonCreator + private BadQueryContextException( + @JsonProperty("error") String errorCode, + @JsonProperty("errorMessage") String errorMessage, + @JsonProperty("errorClass") String errorClass + ) + { + super(errorCode, errorMessage, errorClass); + } +} 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 d8bf22ebc7f..67cb49be915 100644 --- a/processing/src/main/java/org/apache/druid/query/QueryContexts.java +++ b/processing/src/main/java/org/apache/druid/query/QueryContexts.java @@ -418,9 +418,14 @@ public class QueryContexts public static long getTimeout(Query query, long defaultTimeout) { - final long timeout = parseLong(query, TIMEOUT_KEY, defaultTimeout); - Preconditions.checkState(timeout >= 0, "Timeout must be a non negative value, but was [%s]", timeout); - return timeout; + try { + final long timeout = parseLong(query, TIMEOUT_KEY, defaultTimeout); + Preconditions.checkState(timeout >= 0, "Timeout must be a non negative value, but was [%s]", timeout); + return timeout; + } + catch (NumberFormatException e) { + throw new BadQueryContextException(e); + } } public static Query withTimeout(Query query, long timeout) diff --git a/processing/src/test/java/org/apache/druid/query/QueryContextsTest.java b/processing/src/test/java/org/apache/druid/query/QueryContextsTest.java index f5ab7afb264..2ee9b9363e1 100644 --- a/processing/src/test/java/org/apache/druid/query/QueryContextsTest.java +++ b/processing/src/test/java/org/apache/druid/query/QueryContextsTest.java @@ -181,6 +181,21 @@ public class QueryContextsTest QueryContexts.getBrokerServiceName(queryContext); } + @Test + public void testGetTimeout_withNonNumericValue() + { + Map queryContext = new HashMap<>(); + queryContext.put(QueryContexts.TIMEOUT_KEY, "2000'"); + + exception.expect(BadQueryContextException.class); + QueryContexts.getTimeout(new TestQuery( + new TableDataSource("test"), + new MultipleIntervalSegmentSpec(ImmutableList.of(Intervals.of("0/100"))), + false, + queryContext + )); + } + @Test public void testDefaultEnableQueryDebugging() { diff --git a/server/src/main/java/org/apache/druid/server/QueryResource.java b/server/src/main/java/org/apache/druid/server/QueryResource.java index d71fdab56ca..f2a55242ea2 100644 --- a/server/src/main/java/org/apache/druid/server/QueryResource.java +++ b/server/src/main/java/org/apache/druid/server/QueryResource.java @@ -52,7 +52,6 @@ import org.apache.druid.query.QueryInterruptedException; import org.apache.druid.query.QueryTimeoutException; import org.apache.druid.query.QueryToolChest; import org.apache.druid.query.QueryUnsupportedException; -import org.apache.druid.query.ResourceLimitExceededException; import org.apache.druid.query.TruncatedResponseContextException; import org.apache.druid.query.context.ResponseContext; import org.apache.druid.query.context.ResponseContext.Keys; @@ -326,7 +325,7 @@ public class QueryResource implements QueryCountStatsProvider queryLifecycle.emitLogsAndMetrics(unsupported, req.getRemoteAddr(), -1); return ioReaderWriter.getResponseWriter().gotUnsupported(unsupported); } - catch (BadJsonQueryException | ResourceLimitExceededException e) { + catch (BadQueryException e) { interruptedQueryCount.incrementAndGet(); queryLifecycle.emitLogsAndMetrics(e, req.getRemoteAddr(), -1); return ioReaderWriter.getResponseWriter().gotBadQuery(e); 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 a14c9eeaab8..bc3fa82fe46 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 @@ -38,7 +38,6 @@ import org.apache.druid.query.QueryContext; import org.apache.druid.query.QueryInterruptedException; import org.apache.druid.query.QueryTimeoutException; import org.apache.druid.query.QueryUnsupportedException; -import org.apache.druid.query.ResourceLimitExceededException; import org.apache.druid.server.initialization.ServerConfig; import org.apache.druid.server.security.Access; import org.apache.druid.server.security.AuthorizationUtils; @@ -196,7 +195,7 @@ public class SqlResource endLifecycle(sqlQueryId, lifecycle, timeout, remoteAddr, -1); return buildNonOkResponse(QueryTimeoutException.STATUS_CODE, timeout, sqlQueryId); } - catch (SqlPlanningException | ResourceLimitExceededException e) { + catch (BadQueryException e) { endLifecycle(sqlQueryId, lifecycle, e, remoteAddr, -1); return buildNonOkResponse(BadQueryException.STATUS_CODE, e, sqlQueryId); } 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 d7db27ebb65..77a7c7720bc 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 @@ -46,6 +46,7 @@ import org.apache.druid.java.util.common.guava.Sequence; import org.apache.druid.java.util.common.io.Closer; import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.math.expr.ExprMacroTable; +import org.apache.druid.query.BadQueryContextException; import org.apache.druid.query.BaseQuery; import org.apache.druid.query.DefaultQueryConfig; import org.apache.druid.query.Query; @@ -1601,6 +1602,30 @@ public class SqlResourceTest extends CalciteTestBase Assert.assertEquals(Status.OK.getStatusCode(), response.getStatus()); } + @Test + public void testQueryContextException() throws Exception + { + final String sqlQueryId = "badQueryContextTimeout"; + Map queryContext = ImmutableMap.of(QueryContexts.TIMEOUT_KEY, "2000'", BaseQuery.SQL_QUERY_ID, sqlQueryId); + final QueryException queryContextException = doPost( + new SqlQuery( + "SELECT 1337", + ResultFormat.OBJECT, + false, + false, + false, + queryContext, + null + ) + ).lhs; + Assert.assertNotNull(queryContextException); + Assert.assertEquals(BadQueryContextException.ERROR_CODE, queryContextException.getErrorCode()); + Assert.assertEquals(BadQueryContextException.ERROR_CLASS, queryContextException.getErrorClass()); + Assert.assertTrue(queryContextException.getMessage().contains("For input string: \"2000'\"")); + checkSqlRequestLog(false); + Assert.assertTrue(lifecycleManager.getAll(sqlQueryId).isEmpty()); + } + @SuppressWarnings("unchecked") private void checkSqlRequestLog(boolean success) {