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 dcd43d75468d..14eaa4fc5bc6 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 @@ -248,7 +248,7 @@ public void processSqlWithMultiStageQueryEnginePost(String query, @Suspended Asy @Path("query/compare") @ApiOperation(value = "Query Pinot using both the single-stage query engine and the multi-stage query engine and " + "compare the results. The 'sql' field should be set in the request JSON to run the same query on both the " - + "query engines. Set '" + Request.V1SQL + "' and '" + Request.V2SQL + "' if the query needs to be adapted for " + + "query engines. Set '" + Request.SQL_V1 + "' and '" + Request.SQL_V2 + "' if the query needs to be adapted for " + "the two query engines.") @ApiResponses(value = { @ApiResponse(code = 200, message = "Query result comparison response"), @@ -262,15 +262,20 @@ public void processSqlQueryWithBothEnginesAndCompareResults(String query, @Suspe JsonNode requestJson = JsonUtils.stringToJsonNode(query); String v1Query; String v2Query; - if (requestJson.has(Request.SQL)) { - v1Query = requestJson.get(Request.SQL).asText(); - v2Query = v1Query; - } else if (requestJson.has(Request.V1SQL) && requestJson.has(Request.V2SQL)) { - v1Query = requestJson.get(Request.V1SQL).asText(); - v2Query = requestJson.get(Request.V2SQL).asText(); + + if (!requestJson.has(Request.SQL)) { + if (!requestJson.has(Request.SQL_V1) || !requestJson.has(Request.SQL_V2)) { + throw new IllegalStateException("Payload should either contain the query string field '" + Request.SQL + "' " + + "or both of '" + Request.SQL_V1 + "' and '" + Request.SQL_V2 + "'"); + } else { + v1Query = requestJson.get(Request.SQL_V1).asText(); + v2Query = requestJson.get(Request.SQL_V2).asText(); + } } else { - throw new IllegalStateException("Payload should either contain the query string field '" + Request.SQL + "' " - + "or both of '" + Request.V1SQL + "' and '" + Request.V2SQL + "'"); + v1Query = requestJson.has(Request.SQL_V1) ? requestJson.get(Request.SQL_V1).asText() + : requestJson.get(Request.SQL).asText(); + v2Query = requestJson.has(Request.SQL_V2) ? requestJson.get(Request.SQL_V2).asText() + : requestJson.get(Request.SQL).asText(); } ObjectNode v1RequestJson = requestJson.deepCopy(); 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 index b772f3340e8b..1a7c65648139 100644 --- 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 @@ -129,8 +129,8 @@ public void testPinotQueryComparisonApiDifferentQuery() throws Exception { when(request.getRequestURL()).thenReturn(new StringBuilder()); when(_requestHandler.handleRequest(any(), any(), any(), any(), any())) .thenReturn(BrokerResponseNative.EMPTY_RESULT); - _pinotClientRequest.processSqlQueryWithBothEnginesAndCompareResults("{\"v1sql\": \"SELECT v1 FROM mytable\"," - + "\"v2sql\": \"SELECT v2 FROM mytable\"}", asyncResponse, request, null); + _pinotClientRequest.processSqlQueryWithBothEnginesAndCompareResults("{\"sqlV1\": \"SELECT v1 FROM mytable\"," + + "\"sqlV2\": \"SELECT v2 FROM mytable\"}", asyncResponse, request, null); ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(JsonNode.class); verify(_requestHandler, times(2)).handleRequest(requestCaptor.capture(), any(), any(), any(), any()); 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 0ce890e4dc35..2f9b96e33d54 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 @@ -370,8 +370,8 @@ public static class Broker { public static class Request { public static final String SQL = "sql"; - public static final String V1SQL = "v1sql"; - public static final String V2SQL = "v2sql"; + public static final String SQL_V1 = "sqlV1"; + public static final String SQL_V2 = "sqlV2"; public static final String TRACE = "trace"; public static final String QUERY_OPTIONS = "queryOptions";