Add integration test for rebalance in upsert tables (apache#11568)

* Add integration test for rebalance in upsert tables * Add reload test as well * Reload tests * Remove upsert integration test since partial upsert tests cover both * Add explicit status checks for rebalance/reload jobs rather than sleep * Fix flakiness due to segment commit during reload * Use new RebalanceConfig class * Add missing partial upsert cases * Refactor jobstatus to segment reload status --------- Co-authored-by: Kartik Khare <[email protected]> Co-authored-by: Kartik Khare <[email protected]>
soumitra-st · Oct 9, 2023 · 52d16f7 · 52d16f7
1 parent a0964ec
commit 52d16f7
Show file tree

Hide file tree

Showing 5 changed files with 468 additions and 11 deletions.
diff --git a/...t-controller/src/main/java/org/apache/pinot/controller/helix/ControllerRequestClient.java b/...t-controller/src/main/java/org/apache/pinot/controller/helix/ControllerRequestClient.java
@@ -175,11 +175,13 @@ public void resetSegment(String tableNameWithType, String segmentName, String ta
     }
   }
 
-  public void reloadTable(String tableName, TableType tableType, boolean forceDownload)
+  public String reloadTable(String tableName, TableType tableType, boolean forceDownload)
       throws IOException {
     try {
-      HttpClient.wrapAndThrowHttpException(_httpClient.sendJsonPostRequest(new URL(
+      SimpleHttpResponse simpleHttpResponse =
+          HttpClient.wrapAndThrowHttpException(_httpClient.sendJsonPostRequest(new URL(
           _controllerRequestURLBuilder.forTableReload(tableName, tableType, forceDownload)).toURI(), null));
+      return simpleHttpResponse.getResponse();
     } catch (HttpErrorStatusException | URISyntaxException e) {
       throw new IOException(e);
     }

diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/ControllerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/ControllerTest.java
@@ -716,24 +716,24 @@ public long getTableSize(String tableName)
     return getControllerRequestClient().getTableSize(tableName);
   }
 
-  public void reloadOfflineTable(String tableName)
+  public String reloadOfflineTable(String tableName)
       throws IOException {
-    reloadOfflineTable(tableName, false);
+    return reloadOfflineTable(tableName, false);
   }
 
-  public void reloadOfflineTable(String tableName, boolean forceDownload)
+  public String reloadOfflineTable(String tableName, boolean forceDownload)
       throws IOException {
-    getControllerRequestClient().reloadTable(tableName, TableType.OFFLINE, forceDownload);
+    return getControllerRequestClient().reloadTable(tableName, TableType.OFFLINE, forceDownload);
   }
 
   public void reloadOfflineSegment(String tableName, String segmentName, boolean forceDownload)
       throws IOException {
     getControllerRequestClient().reloadSegment(tableName, segmentName, forceDownload);
   }
 
-  public void reloadRealtimeTable(String tableName)
+  public String reloadRealtimeTable(String tableName)
       throws IOException {
-    getControllerRequestClient().reloadTable(tableName, TableType.REALTIME, false);
+    return getControllerRequestClient().reloadTable(tableName, TableType.REALTIME, false);
   }
 
   public void createBrokerTenant(String tenantName, int numBrokers)

diff --git a/...tests/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTestSet.java b/...tests/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTestSet.java
@@ -651,7 +651,7 @@ public String reloadTableAndValidateResponse(String tableName, TableType tableTy
     String isZKWriteSuccess = tableLevelDetails.get("reloadJobMetaZKStorageStatus").asText();
     assertEquals(isZKWriteSuccess, "SUCCESS");
     String jobId = tableLevelDetails.get("reloadJobId").asText();
-    String jobStatusResponse = sendGetRequest(_controllerRequestURLBuilder.forControllerJobStatus(jobId));
+    String jobStatusResponse = sendGetRequest(_controllerRequestURLBuilder.forSegmentReloadStatus(jobId));
     JsonNode jobStatus = JsonUtils.stringToJsonNode(jobStatusResponse);
 
     // Validate all fields are present
@@ -663,7 +663,7 @@ public String reloadTableAndValidateResponse(String tableName, TableType tableTy
 
   public boolean isReloadJobCompleted(String reloadJobId)
       throws Exception {
-    String jobStatusResponse = sendGetRequest(_controllerRequestURLBuilder.forControllerJobStatus(reloadJobId));
+    String jobStatusResponse = sendGetRequest(_controllerRequestURLBuilder.forSegmentReloadStatus(reloadJobId));
     JsonNode jobStatus = JsonUtils.stringToJsonNode(jobStatusResponse);
 
     assertEquals(jobStatus.get("metadata").get("jobId").asText(), reloadJobId);