Added pinot-error-code header in query response (apache#12338)

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();
+  }
 }

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);
+  }
+}

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";