yashmayya commented on code in PR #15277: URL: https://github.com/apache/pinot/pull/15277#discussion_r2003003191
########## pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/QueryAssert.java: ########## @@ -0,0 +1,159 @@ +/** + * 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.pinot.integration.tests; + +import com.fasterxml.jackson.databind.JsonNode; +import java.util.Locale; +import org.apache.pinot.spi.exception.QueryErrorCode; +import org.assertj.core.api.AbstractAssert; +import org.assertj.core.api.AbstractSoftAssertions; + +/// A custom AssertJ assertion class for query responses that provides a fluent API for asserting on query responses. +/// +/// The current implementation is partial and we should be adding more methods to support more use cases as more tests +/// are migrated to this class instead of TestNG's Assert class. +public class QueryAssert extends AbstractAssert<QueryAssert, JsonNode> { + public QueryAssert(JsonNode jsonNode) { + super(jsonNode, QueryAssert.class); + } + + public static QueryAssert assertThat(JsonNode actual) { + return new QueryAssert(actual); + } + + public QueryAssert hasNoExceptions() { + isNotNull(); + if (actual.has("exceptions") && !actual.get("exceptions").isEmpty()) { + failWithMessage("Expected no exceptions but found <%s>", actual.get("exceptions")); + } + return this; + } + + /// Obtains the first exception in the query response, returning it as a [QueryErrorAssert] object. + /// + /// It fails if there are no exceptions in the query response. + public QueryErrorAssert firstException() { + isNotNull(); + if (!actual.has("exceptions")) { + failWithMessage("No exceptions found in query response"); + } + JsonNode exceptions = actual.get("exceptions"); + if (exceptions.isEmpty()) { + failWithMessage("No exceptions found in query response"); + } + return new QueryErrorAssert(actual.get("exceptions").get(0)); + } + + + /// Obtains the first exception in the query response, returning it as a [QueryErrorAssert.Soft] object. + /// + /// It fails if there are no exceptions in the query response. + /// + /// Unlike [#firstException], this method returns a [QueryErrorAssert.Soft] object, which allows for multiple + /// assertions to be made before failing. + /// + /// See [Soft Assertions in AssertJ docs](https://assertj.github.io/doc/#assertj-core-soft-assertions) + public QueryErrorAssert.Soft softFirstException() { Review Comment: When would we explicitly want non-soft assertions? ########## pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/MultiStageEngineIntegrationTest.java: ########## @@ -703,7 +705,9 @@ public void testBase64Func() // invalid argument sqlQuery = "SELECT fromBase64('hello!') FROM mytable"; response = postQuery(sqlQuery); - assertTrue(response.get("exceptions").get(0).get("message").toString().contains("Illegal base64 character")); + assertThat(response.get("exceptions").get(0).get("message").asText()) + .as("First exception message") + .contains("Illegal base64 character"); Review Comment: Why not use `QueryAssert` here? ########## pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java: ########## @@ -192,17 +195,8 @@ private WorkerManager getWorkerManager(SqlNodeAndOptions sqlNodeAndOptions) { * @return QueryPlannerResult containing the dispatchable query plan and the relRoot. */ public QueryPlannerResult planQuery(String sqlQuery, SqlNodeAndOptions sqlNodeAndOptions, long requestId) { - try (PlannerContext plannerContext = getPlannerContext(sqlNodeAndOptions)) { - RelRoot relRoot = compileQuery(sqlNodeAndOptions.getSqlNode(), plannerContext); - // TODO: current code only assume one SubPlan per query, but we should support multiple SubPlans per query. - // Each SubPlan should be able to run independently from Broker then set the results into the dependent - // SubPlan for further processing. - DispatchableSubPlan dispatchableSubPlan = toDispatchableSubPlan(relRoot, plannerContext, requestId); - return getQueryPlannerResult(plannerContext, dispatchableSubPlan, null, dispatchableSubPlan.getTableNames()); - } catch (CalciteContextException e) { - throw new RuntimeException("Error composing query plan for '" + sqlQuery + "': " + e.getMessage() + "'", e); - } catch (Throwable t) { - throw new RuntimeException("Error composing query plan for: " + sqlQuery, t); + try (CompiledQuery compiledQuery = compile(sqlQuery, sqlNodeAndOptions)) { + return compiledQuery.planQuery(requestId); Review Comment: Neat refactor! ########## pinot-common/src/main/java/org/apache/pinot/common/utils/Timer.java: ########## @@ -18,32 +18,34 @@ */ package org.apache.pinot.common.utils; +import java.util.concurrent.TimeUnit; + /** * Utility class that works with a timeout in milliseconds and provides methods to check remaining time and expiration. */ public class Timer { - private final long _timeoutMillis; - private final long _startTime; + private final long _startNs; + private final long _deadlineNs; + private final ClockNs _clock; - /** - * Initializes the Timer with the specified timeout in milliseconds. - * - * @param timeoutMillis the timeout duration in milliseconds - */ - public Timer(long timeoutMillis) { - _timeoutMillis = timeoutMillis; - _startTime = System.currentTimeMillis(); + public Timer(Long timeout, TimeUnit timeUnit) { + this(System::nanoTime, timeout, timeUnit); + } + + public Timer(ClockNs clock, Long timeout, TimeUnit timeUnit) { + _clock = clock; + _startNs = _clock.nanos(); + _deadlineNs = timeUnit.toNanos(timeout) + _clock.nanos(); } /** * Returns the remaining time in milliseconds. If the timeout has expired, it returns 0. * * @return the remaining time in milliseconds */ - public long getRemainingTime() { - long elapsedTime = System.currentTimeMillis() - _startTime; - long remainingTime = _timeoutMillis - elapsedTime; - return Math.max(remainingTime, 0); + public long getRemainingTimeMs() { + long remainingNs = _deadlineNs - _clock.nanos(); + return Math.max(remainingNs / 1000, 0); Review Comment: Shouldn't this be `remainingNs / 1000000`? ########## pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/QueryAssert.java: ########## @@ -0,0 +1,159 @@ +/** + * 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.pinot.integration.tests; + +import com.fasterxml.jackson.databind.JsonNode; +import java.util.Locale; +import org.apache.pinot.spi.exception.QueryErrorCode; +import org.assertj.core.api.AbstractAssert; +import org.assertj.core.api.AbstractSoftAssertions; + +/// A custom AssertJ assertion class for query responses that provides a fluent API for asserting on query responses. +/// +/// The current implementation is partial and we should be adding more methods to support more use cases as more tests +/// are migrated to this class instead of TestNG's Assert class. +public class QueryAssert extends AbstractAssert<QueryAssert, JsonNode> { + public QueryAssert(JsonNode jsonNode) { + super(jsonNode, QueryAssert.class); + } + + public static QueryAssert assertThat(JsonNode actual) { + return new QueryAssert(actual); + } + + public QueryAssert hasNoExceptions() { + isNotNull(); + if (actual.has("exceptions") && !actual.get("exceptions").isEmpty()) { + failWithMessage("Expected no exceptions but found <%s>", actual.get("exceptions")); + } + return this; + } + + /// Obtains the first exception in the query response, returning it as a [QueryErrorAssert] object. + /// + /// It fails if there are no exceptions in the query response. + public QueryErrorAssert firstException() { + isNotNull(); + if (!actual.has("exceptions")) { + failWithMessage("No exceptions found in query response"); + } + JsonNode exceptions = actual.get("exceptions"); + if (exceptions.isEmpty()) { + failWithMessage("No exceptions found in query response"); + } + return new QueryErrorAssert(actual.get("exceptions").get(0)); + } + + + /// Obtains the first exception in the query response, returning it as a [QueryErrorAssert.Soft] object. + /// + /// It fails if there are no exceptions in the query response. + /// + /// Unlike [#firstException], this method returns a [QueryErrorAssert.Soft] object, which allows for multiple + /// assertions to be made before failing. + /// + /// See [Soft Assertions in AssertJ docs](https://assertj.github.io/doc/#assertj-core-soft-assertions) + public QueryErrorAssert.Soft softFirstException() { + isNotNull(); + if (!actual.has("exceptions")) { + failWithMessage("No exceptions found in query response"); + } + JsonNode exceptions = actual.get("exceptions"); + if (exceptions.isEmpty()) { + failWithMessage("No exceptions found in query response"); + } + + return new QueryErrorAssert.Soft(actual.get("exceptions").get(0)); + } + + public static class Soft extends AbstractSoftAssertions implements AutoCloseable { Review Comment: Unused? ########## pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java: ########## @@ -278,53 +204,69 @@ private QueryEnvironment.QueryPlannerResult getQueryPlannerResult(PlannerContext return new QueryPlannerResult(dispatchableSubPlan, explainStr, tableNames, extraFields); } + /// @deprecated Use [#compile] and then [explain][CompiledQuery#explain(long) ] the returned query instead @VisibleForTesting + @Deprecated public String explainQuery(String sqlQuery, long requestId) { SqlNodeAndOptions sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions(sqlQuery); - QueryPlannerResult queryPlannerResult = explainQuery(sqlQuery, sqlNodeAndOptions, requestId, null); - return queryPlannerResult.getExplainPlan(); + try (CompiledQuery compiledQuery = compile(sqlQuery, sqlNodeAndOptions)) { + QueryPlannerResult queryPlannerResult = compiledQuery.explain(requestId, null); + return queryPlannerResult.getExplainPlan(); + } } - public List<String> getTableNamesForQuery(String sqlQuery) { - SqlNodeAndOptions sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions(sqlQuery); - try (PlannerContext plannerContext = getPlannerContext(sqlNodeAndOptions)) { + public CompiledQuery compile(String sqlQuery) { + return compile(sqlQuery, CalciteSqlParser.compileToSqlNodeAndOptions(sqlQuery)); + } + + /// Given a query, parses, validates and optimizes the query into a [CompiledQuery]. + /// + /// The returned query can then be planned, explained or used to get the tables involved in the query. + /// + /// @throws QueryException if the query cannot be compiled. Usual error types are [QueryErrorCode#SQL_PARSING] and + /// [QueryErrorCode#QUERY_VALIDATION]. [QueryErrorCode#QUERY_EXECUTION] is also possible if there is an error when + /// a function call is reduced into a constant. Review Comment: Is that the only possible source of `QUERY_EXECUTION` error here? ########## pinot-query-planner/src/main/java/org/apache/pinot/query/CalciteContextExceptionClassifier.java: ########## @@ -0,0 +1,50 @@ +/** + * 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.pinot.query; + +import java.util.regex.Pattern; +import org.apache.calcite.runtime.CalciteContextException; +import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.exception.QueryException; + + +public class CalciteContextExceptionClassifier { + + private static final Pattern UNKNOWN_COLUMN_PATTERN = Pattern.compile(".*Column '.+' not found in any table"); + private static final Pattern UNKNOWN_TABLE_PATTERN = Pattern.compile(".*Object '.+' not found"); + + private CalciteContextExceptionClassifier() { + } + + /// Analyzes the exception and classifies it as a [QueryErrorCode] if possible. + /// + /// Returns a [QueryException#QUERY_VALIDATION] exception if the exception is not recognized. Review Comment: `QueryException` -> `QueryErrorCode`? ########## pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java: ########## @@ -423,25 +504,32 @@ private static void throwTableAccessError(TableAuthorizationResult tableAuthoriz } /** - * Runs the query planning in a separate thread so that we can enforce a timeout on it (in some rare cases, - * we can see query compilation taking a very long time). + * Calls the given callable in a separate thread and enforces a timeout on it. + * + * The only exception that can be thrown by this method is a QueryException. All other exceptions are caught and + * wrapped in a QueryException. Specifically, {@link TimeoutException} is caught and wrapped in a QueryException with + * the error code {@link QueryErrorCode#BROKER_TIMEOUT} and other exceptions are treated as internal errors.K Review Comment: ```suggestion * the error code {@link QueryErrorCode#BROKER_TIMEOUT} and other exceptions are treated as internal errors. ``` ########## pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/ErrorCodesIntegrationTest.java: ########## @@ -0,0 +1,215 @@ +/** + * 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.pinot.integration.tests; + +import java.io.File; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.commons.io.FileUtils; +import org.apache.helix.model.HelixConfigScope; +import org.apache.helix.model.builder.HelixConfigScopeBuilder; +import org.apache.pinot.spi.config.table.FieldConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.util.TestUtils; +import org.intellij.lang.annotations.Language; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertTrue; + + +public abstract class ErrorCodesIntegrationTest extends BaseClusterIntegrationTestSet { + private static final int NUM_BROKERS = 1; + private static final int NUM_SERVERS = 1; + private static final int NUM_SEGMENTS = 1; Review Comment: Unused? ########## pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java: ########## @@ -531,4 +619,20 @@ public FailureDetector.ServerState retryUnhealthyServer(String instanceId) { return _queryDispatcher.checkConnectivityToInstance(serverInstance); } + + public static boolean isYellowError(QueryException e) { + switch (e.getErrorCode()) { + case QUERY_SCHEDULING_TIMEOUT: + case EXECUTION_TIMEOUT: + case INTERNAL: + case UNKNOWN: Review Comment: Why are these "yellow" errors and not "red"? ########## pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java: ########## @@ -558,4 +546,134 @@ default boolean defaultEnableGroupTrim() { @Nullable WorkerManager getWorkerManager(); } + + /// A query that have been parsed, validates, transformed into a [RelNode] and optimized with Calcite. + /// + /// This represents the last point where Calcite is being used. This object can then be: + /// - Used to get the tables involved in the query (see [#getTableNames]) + /// - Used to explain the query plan (see [#explain]) + /// - Used to plan how to evaluate the query using Pinot Engine (see [#planQuery]) + /// + /// Compiled queries are craeted by calling [QueryEnvironment#compile] and should be closed. Review Comment: Let's also document that closing the compiled query closes the passed in `PlannerContext`? ########## pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java: ########## @@ -278,53 +204,69 @@ private QueryEnvironment.QueryPlannerResult getQueryPlannerResult(PlannerContext return new QueryPlannerResult(dispatchableSubPlan, explainStr, tableNames, extraFields); } + /// @deprecated Use [#compile] and then [explain][CompiledQuery#explain(long) ] the returned query instead @VisibleForTesting + @Deprecated public String explainQuery(String sqlQuery, long requestId) { SqlNodeAndOptions sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions(sqlQuery); - QueryPlannerResult queryPlannerResult = explainQuery(sqlQuery, sqlNodeAndOptions, requestId, null); - return queryPlannerResult.getExplainPlan(); + try (CompiledQuery compiledQuery = compile(sqlQuery, sqlNodeAndOptions)) { + QueryPlannerResult queryPlannerResult = compiledQuery.explain(requestId, null); + return queryPlannerResult.getExplainPlan(); + } } - public List<String> getTableNamesForQuery(String sqlQuery) { - SqlNodeAndOptions sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions(sqlQuery); - try (PlannerContext plannerContext = getPlannerContext(sqlNodeAndOptions)) { + public CompiledQuery compile(String sqlQuery) { + return compile(sqlQuery, CalciteSqlParser.compileToSqlNodeAndOptions(sqlQuery)); + } + + /// Given a query, parses, validates and optimizes the query into a [CompiledQuery]. + /// + /// The returned query can then be planned, explained or used to get the tables involved in the query. + /// + /// @throws QueryException if the query cannot be compiled. Usual error types are [QueryErrorCode#SQL_PARSING] and + /// [QueryErrorCode#QUERY_VALIDATION]. [QueryErrorCode#QUERY_EXECUTION] is also possible if there is an error when + /// a function call is reduced into a constant. Review Comment: I don't know if it's just an issue with my IDE but these query error codes don't render in the Javadoc ########## pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java: ########## @@ -151,8 +161,13 @@ private PlannerContext getPlannerContext(SqlNodeAndOptions sqlNodeAndOptions) { sqlNodeAndOptions.getOptions(), _envConfig, format); } - public Set<String> getResolvedTables() { - return _catalog.getResolvedTables(); + /// @deprecated Use [#compile] and then [plan][CompiledQuery#planQuery(long) ] the returned query instead Review Comment: ```suggestion /// @deprecated Use [#compile] and then [plan][CompiledQuery#planQuery(long)] the returned query instead ``` nit: the extra whitespace causes the closing bracket to actually be rendered ########## pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotQueryResource.java: ########## @@ -105,47 +104,41 @@ public class PinotQueryResource { @POST @Path("sql") @ManualAuthorization // performed by broker - public String handlePostSql(String requestJsonStr, @Context HttpHeaders httpHeaders) { + public StreamingOutput handlePostSql(String requestJsonStr, @Context HttpHeaders httpHeaders) { Review Comment: What's the purpose of changing the return type to `StreamingOutput`? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org For additional commands, e-mail: commits-h...@pinot.apache.org
