Add a register endpoint to upload zk metadata directly

pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java

@@ -71,6 +71,9 @@
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.tuple.ImmutablePair;
 import org.apache.commons.lang3.tuple.Pair;
+import org.apache.helix.model.IdealState;
+import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.apache.pinot.common.metadata.segment.SegmentZKMetadata;
 import org.apache.pinot.common.metrics.ControllerGauge;
 import org.apache.pinot.common.metrics.ControllerMeter;
 import org.apache.pinot.common.metrics.ControllerMetrics;
@@ -986,6 +989,135 @@ public void uploadSegmentAsMultiPartV2(FormDataMultiPart multiPart,
     }
   }
 
+  @POST
+  @ManagedAsync
+  @Produces(MediaType.APPLICATION_JSON)
+  @Consumes(MediaType.APPLICATION_JSON)
+  @Path("/segments/register")
+  @Authorize(targetType = TargetType.TABLE, paramName = "tableName", action = Actions.Table.UPLOAD_SEGMENT)
+  @Authenticate(AccessType.CREATE)
+  @ApiOperation(value = "Register a segment from pre-built ZK metadata",
+      notes = "Registers a segment directly from its ZK metadata JSON without file upload, unpacking, or copying. "
+          + "The segment files must already exist at the downloadUrl in the ZK metadata. "
+          + "Request body is SegmentZKMetadata JSON: {\"segmentName\": \"...\", \"simpleFields\": {...}}. "
+          + "Required simpleFields: segment.download.url, segment.crc. "
+          + "Optional but recommended: segment.total.docs, segment.start.time, segment.end.time, "
+          + "segment.time.unit, segment.index.version, segment.creation.time.")
+  @ApiResponses(value = {
+      @ApiResponse(code = 200, message = "Successfully registered segment"),
+      @ApiResponse(code = 400, message = "Bad Request"),
+      @ApiResponse(code = 409, message = "Segment already exists and refresh not permitted"),
+      @ApiResponse(code = 500, message = "Internal error")
+  })
+  @TrackInflightRequestMetrics
+  @TrackedByGauge(gauge = ControllerGauge.SEGMENT_UPLOADS_IN_PROGRESS)
+  public void registerSegment(
+      String segmentZKMetadataJson,
+      @ApiParam(value = "Name of the table", required = true)
+      @QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_NAME) String tableName,
+      @ApiParam(value = "Type of the table")
+      @QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_TYPE) @DefaultValue("OFFLINE") String tableType,
+      @ApiParam(value = "Whether to refresh if the segment already exists")
+      @QueryParam(FileUploadDownloadClient.QueryParameters.ALLOW_REFRESH) @DefaultValue("true") boolean allowRefresh,
+      @Suspended final AsyncResponse asyncResponse) {
+    try {
+      asyncResponse.resume(registerSegmentInternal(segmentZKMetadataJson, tableName, tableType, allowRefresh));
+    } catch (Throwable t) {
+      asyncResponse.resume(t);
+    }
+  }
+
+  private SuccessResponse registerSegmentInternal(String segmentZKMetadataJson, String tableName, String tableType,
+      boolean allowRefresh) {
+    if (StringUtils.isEmpty(tableName)) {
+      throw new ControllerApplicationException(LOGGER, "tableName is required", Response.Status.BAD_REQUEST);
+    }
+
+    SegmentZKMetadata segmentZKMetadata;
+    try {
+      segmentZKMetadata = SegmentZKMetadata.fromJsonString(segmentZKMetadataJson);
+    } catch (Exception e) {
+      throw new ControllerApplicationException(LOGGER,
+          "Failed to parse segment ZK metadata: " + e.getMessage(), Response.Status.BAD_REQUEST);
+    }
+
+    String segmentName = segmentZKMetadata.getSegmentName();
+    if (StringUtils.isEmpty(segmentName)) {
+      throw new ControllerApplicationException(LOGGER, "segmentName is required in ZK metadata",
+          Response.Status.BAD_REQUEST);
+    }
+    if (StringUtils.isEmpty(segmentZKMetadata.getDownloadUrl())) {
+      throw new ControllerApplicationException(LOGGER,
+          "segment.download.url is required in simpleFields", Response.Status.BAD_REQUEST);
+    }
+
+    String rawTableName = TableNameBuilder.extractRawTableName(tableName);
+    TableType type = TableType.valueOf(tableType.toUpperCase());
+    String tableNameWithType = TableNameBuilder.forType(type).tableNameWithType(rawTableName);
+
+    TableConfig tableConfig = _pinotHelixResourceManager.getTableConfig(tableNameWithType);
+    if (tableConfig == null) {
+      throw new ControllerApplicationException(LOGGER, "Failed to find table: " + tableNameWithType,
+          Response.Status.BAD_REQUEST);
+    }
+
+    LOGGER.info("Registering segment: {} for table: {}", segmentName, tableNameWithType);
+
+    ZNRecord existingZNRecord = _pinotHelixResourceManager.getSegmentMetadataZnRecord(tableNameWithType, segmentName);
+
+    // If ZK metadata exists but the segment is absent from the ideal state, a previous registration failed
+    // partway through. Remove the stale ZK entry and treat this as a new segment.
+    if (existingZNRecord != null) {
+      IdealState idealState = _pinotHelixResourceManager.getTableIdealState(tableNameWithType);
+      Preconditions.checkState(idealState != null, "Failed to find ideal state for table: %s", tableNameWithType);
+      if (idealState.getInstanceStateMap(segmentName) == null) {
+        LOGGER.warn("Removing stale segment ZK metadata (recovering from previous registration failure) "
+            + "for table: {}, segment: {}", tableNameWithType, segmentName);
+        Preconditions.checkState(_pinotHelixResourceManager.removeSegmentZKMetadata(tableNameWithType, segmentName),
+            "Failed to remove stale segment ZK metadata for table: %s, segment: %s", tableNameWithType, segmentName);
+        existingZNRecord = null;
+      }
+    }
+
+    if (existingZNRecord != null) {
+      if (!allowRefresh) {
+        throw new ControllerApplicationException(LOGGER,
+            String.format("Segment: %s already exists in table: %s. Refresh not permitted.", segmentName,
+                tableNameWithType), Response.Status.CONFLICT);
+      }
+      LOGGER.info("Segment: {} already exists in table: {}, refreshing", segmentName, tableNameWithType);
+      segmentZKMetadata.setRefreshTime(System.currentTimeMillis());
+      if (!_pinotHelixResourceManager.updateZkMetadata(tableNameWithType, segmentZKMetadata,
+          existingZNRecord.getVersion())) {
+        throw new RuntimeException(
+            String.format("Failed to update ZK metadata for segment: %s of table: %s", segmentName,
+                tableNameWithType));
+      }
+      _pinotHelixResourceManager.sendSegmentRefreshMessage(tableNameWithType, segmentName, true, true);
+    } else {
+      if (_pinotHelixResourceManager.needTieredSegmentAssignment(tableConfig)) {
+        _pinotHelixResourceManager.updateSegmentTargetTier(tableNameWithType, segmentZKMetadata,
+            _pinotHelixResourceManager.getSortedTiers(tableConfig));
+      }
+      segmentZKMetadata.setPushTime(System.currentTimeMillis());
+      if (!_pinotHelixResourceManager.createSegmentZkMetadata(tableNameWithType, segmentZKMetadata)) {
+        throw new RuntimeException(
+            String.format("Failed to create ZK metadata for segment: %s of table: %s", segmentName,
+                tableNameWithType));
+      }
+      try {
+        _pinotHelixResourceManager.assignSegment(tableConfig, segmentZKMetadata);
+      } catch (Exception e) {
+        LOGGER.error("Failed to assign segment: {} for table: {}, ZK metadata already cleaned up by assignSegment",
+            segmentName, tableNameWithType, e);
+        throw e;
+      }
+    }
+
+    LOGGER.info("Successfully registered segment: {} for table: {}", segmentName, tableNameWithType);
+    return new SuccessResponse("Successfully registered segment: " + segmentName + " of table: " + tableNameWithType);
+  }
+
   @POST
   @Path("segments/{tableName}/startReplaceSegments")
   @Authorize(targetType = TargetType.TABLE, paramName = "tableName", action = Actions.Table.REPLACE_SEGMENT)