From ac13a191b945a80084f0a2794391e4be2f463252 Mon Sep 17 00:00:00 2001 From: Abhishek Sharma Date: Wed, 28 Feb 2024 20:02:49 -0500 Subject: [PATCH] Added pinot-error-code header in query response (#12338) --- .../api/resources/PinotClientRequest.java | 36 +++++++++++-- .../api/resources/PinotClientRequestTest.java | 54 +++++++++++++++++++ .../pinot/spi/utils/CommonConstants.java | 1 + 3 files changed, 87 insertions(+), 4 deletions(-) create mode 100644 pinot-broker/src/test/java/org/apache/pinot/broker/api/resources/PinotClientRequestTest.java diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotClientRequest.java b/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotClientRequest.java index 989bf6b580b..0bdec44a090 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotClientRequest.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotClientRequest.java @@ -20,6 +20,7 @@ import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ArrayListMultimap; import com.google.common.collect.ImmutableMap; import com.google.common.collect.Multimap; @@ -75,6 +76,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import static org.apache.pinot.spi.utils.CommonConstants.Controller.PINOT_QUERY_ERROR_CODE_HEADER; import static org.apache.pinot.spi.utils.CommonConstants.SWAGGER_AUTHORIZATION_KEY; @@ -125,7 +127,7 @@ public void processSqlQueryGet(@ApiParam(value = "Query", required = true) @Quer requestJson.put(Request.DEBUG_OPTIONS, debugOptions); } BrokerResponse brokerResponse = executeSqlQuery(requestJson, makeHttpIdentity(requestContext), true, httpHeaders); - asyncResponse.resume(brokerResponse.toJsonString()); + asyncResponse.resume(getPinotQueryResponse(brokerResponse)); } catch (WebApplicationException wae) { asyncResponse.resume(wae); } catch (Exception e) { @@ -155,7 +157,7 @@ public void processSqlQueryPost(String query, @Suspended AsyncResponse asyncResp } BrokerResponse brokerResponse = executeSqlQuery((ObjectNode) requestJson, makeHttpIdentity(requestContext), false, httpHeaders); - asyncResponse.resume(brokerResponse.toJsonString()); + asyncResponse.resume(getPinotQueryResponse(brokerResponse)); } catch (WebApplicationException wae) { asyncResponse.resume(wae); } catch (Exception e) { @@ -189,7 +191,7 @@ public void processSqlWithMultiStageQueryEngineGet( requestJson.put(Request.SQL, query); BrokerResponse brokerResponse = executeSqlQuery(requestJson, makeHttpIdentity(requestContext), true, httpHeaders, true); - asyncResponse.resume(brokerResponse.toJsonString()); + asyncResponse.resume(getPinotQueryResponse(brokerResponse)); } catch (WebApplicationException wae) { asyncResponse.resume(wae); } catch (Exception e) { @@ -219,7 +221,7 @@ public void processSqlWithMultiStageQueryEnginePost(String query, @Suspended Asy } BrokerResponse brokerResponse = executeSqlQuery((ObjectNode) requestJson, makeHttpIdentity(requestContext), false, httpHeaders, true); - asyncResponse.resume(brokerResponse.toJsonString()); + asyncResponse.resume(getPinotQueryResponse(brokerResponse)); } catch (WebApplicationException wae) { asyncResponse.resume(wae); } catch (Exception e) { @@ -348,4 +350,30 @@ private static HttpRequesterIdentity makeHttpIdentity(org.glassfish.grizzly.http return identity; } + + /** + * Generate Response object from the BrokerResponse object with 'X-Pinot-Error-Code' header value + * + * If the query is successful the 'X-Pinot-Error-Code' header value is set to -1 + * otherwise, the first error code of the broker response exception array will become the header value + * + * @param brokerResponse + * @return Response + * @throws Exception + */ + @VisibleForTesting + static Response getPinotQueryResponse(BrokerResponse brokerResponse) + throws Exception { + int queryErrorCodeHeaderValue = -1; // default value of the header. + + if (brokerResponse.getExceptionsSize() != 0) { + // set the header value as first exception error code value. + queryErrorCodeHeaderValue = brokerResponse.getProcessingExceptions().get(0).getErrorCode(); + } + + // returning the Response with OK status and header value. + return Response.ok() + .header(PINOT_QUERY_ERROR_CODE_HEADER, queryErrorCodeHeaderValue) + .entity(brokerResponse.toJsonString()).build(); + } } diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/api/resources/PinotClientRequestTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/api/resources/PinotClientRequestTest.java new file mode 100644 index 00000000000..e79b20c57b7 --- /dev/null +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/api/resources/PinotClientRequestTest.java @@ -0,0 +1,54 @@ +/** + * 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.broker.api.resources; + +import javax.ws.rs.core.Response; +import org.apache.pinot.common.response.BrokerResponse; +import org.apache.pinot.common.response.broker.BrokerResponseNative; +import org.testng.Assert; +import org.testng.annotations.Test; + +import static org.apache.pinot.common.exception.QueryException.TABLE_DOES_NOT_EXIST_ERROR_CODE; +import static org.apache.pinot.spi.utils.CommonConstants.Controller.PINOT_QUERY_ERROR_CODE_HEADER; + + +public class PinotClientRequestTest { + + @Test + public void testGetPinotQueryResponse() + throws Exception { + + // for successful query result the 'X-Pinot-Error-Code' should be -1 + BrokerResponse emptyResultBrokerResponse = BrokerResponseNative.EMPTY_RESULT; + Response successfulResponse = PinotClientRequest.getPinotQueryResponse(emptyResultBrokerResponse); + Assert.assertEquals(successfulResponse.getStatus(), Response.Status.OK.getStatusCode()); + Assert.assertTrue(successfulResponse.getHeaders().containsKey(PINOT_QUERY_ERROR_CODE_HEADER)); + Assert.assertEquals(successfulResponse.getHeaders().get(PINOT_QUERY_ERROR_CODE_HEADER).size(), 1); + Assert.assertEquals(successfulResponse.getHeaders().get(PINOT_QUERY_ERROR_CODE_HEADER).get(0), -1); + + // for failed query result the 'X-Pinot-Error-Code' should be Error code fo exception. + BrokerResponse tableDoesNotExistBrokerResponse = BrokerResponseNative.TABLE_DOES_NOT_EXIST; + Response tableDoesNotExistResponse = PinotClientRequest.getPinotQueryResponse(tableDoesNotExistBrokerResponse); + Assert.assertEquals(tableDoesNotExistResponse.getStatus(), Response.Status.OK.getStatusCode()); + Assert.assertTrue(tableDoesNotExistResponse.getHeaders().containsKey(PINOT_QUERY_ERROR_CODE_HEADER)); + Assert.assertEquals(tableDoesNotExistResponse.getHeaders().get(PINOT_QUERY_ERROR_CODE_HEADER).size(), 1); + Assert.assertEquals(tableDoesNotExistResponse.getHeaders().get(PINOT_QUERY_ERROR_CODE_HEADER).get(0), + TABLE_DOES_NOT_EXIST_ERROR_CODE); + } +} diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index 3131e2b42b2..dfe625bfef6 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -765,6 +765,7 @@ public static class Controller { public static final String VERSION_HTTP_HEADER = "Pinot-Controller-Version"; public static final String SEGMENT_NAME_HTTP_HEADER = "Pinot-Segment-Name"; public static final String TABLE_NAME_HTTP_HEADER = "Pinot-Table-Name"; + public static final String PINOT_QUERY_ERROR_CODE_HEADER = "X-Pinot-Error-Code"; public static final String INGESTION_DESCRIPTOR = "Pinot-Ingestion-Descriptor"; public static final String PREFIX_OF_CONFIG_OF_PINOT_CRYPTER = "pinot.controller.crypter";