Rename v1 / v2 sql query fields; allow fallback to 'sql' field

apache · Sep 24, 2024 · 269f934 · 269f934
1 parent cbf303a
commit 269f934
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 13 deletions.
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
@@ -292,7 +292,7 @@ public void processTimeSeriesInstantQuery(@Suspended AsyncResponse asyncResponse
   @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"),
@@ -306,15 +306,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
@@ -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<JsonNode> 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
@@ -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";