Implement registerSegment API to support corp to prod migration

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

@@ -71,6 +71,8 @@
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.tuple.ImmutablePair;
 import org.apache.commons.lang3.tuple.Pair;
+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 +988,112 @@ 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 = 404, message = "Segment file not found at the download URL"),
+      @ApiResponse(code = 409, message = "Segment already exists"),
+      @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,
+      @Suspended final AsyncResponse asyncResponse) {
+    try {
+      asyncResponse.resume(registerSegmentInternal(segmentZKMetadataJson, tableName, tableType));
+    } catch (Throwable t) {
+      asyncResponse.resume(t);
+    }
+  }
+
+  private SuccessResponse registerSegmentInternal(String segmentZKMetadataJson, String tableName, String tableType) {
+    if (StringUtils.isEmpty(tableName)) {
+      throw new ControllerApplicationException(LOGGER, "tableName is required", Response.Status.BAD_REQUEST);
+    }
+
+    SegmentZKMetadata segmentZKMetadata;
+    try {
+      ZNRecord znRecord = JsonUtils.stringToObject(segmentZKMetadataJson, ZNRecord.class);
+      segmentZKMetadata = new SegmentZKMetadata(znRecord);
+    } 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);
+    }
+    String downloadUrl = segmentZKMetadata.getDownloadUrl();
+    if (StringUtils.isEmpty(downloadUrl)) {
+      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);
+
+    try {
+      LOGGER.info("Registering segment: {} for table: {}", segmentName, tableNameWithType);
+
+      TableConfig tableConfig = _pinotHelixResourceManager.getTableConfig(tableNameWithType);
+      if (tableConfig == null) {
+        throw new ControllerApplicationException(LOGGER, "Failed to find table: " + tableNameWithType,
+            Response.Status.BAD_REQUEST);
+      }
+
+      // Validate that the segment file exists in the deep store.
+      // The download URL may be a VIP URL (https://...) while the actual file resides on NFS (file://...).
+      // Resolve the actual deep store path using the controller's data directory.
+      URI dataDirURI = ControllerFilePathProvider.getInstance().getDataDirURI();
+      String encodedSegmentName = URIUtils.encode(segmentName);
+      URI segmentFileURI = URIUtils.getUri(dataDirURI.toString(), rawTableName, encodedSegmentName);
+      PinotFS pinotFS = PinotFSFactory.create(dataDirURI.getScheme());
+      if (!pinotFS.exists(segmentFileURI)) {
+        LOGGER.error("Segment file not found in deep store at: {} for segment: {} of table: {}",
+            segmentFileURI, segmentName, tableNameWithType);
+        throw new ControllerApplicationException(LOGGER,
+            String.format("Segment file not found in deep store at: %s", segmentFileURI),
+            Response.Status.NOT_FOUND);
+      }
+
+      ZKOperator zkOperator = new ZKOperator(_pinotHelixResourceManager, _controllerConf, _controllerMetrics);
+      zkOperator.registerSegment(tableConfig, segmentZKMetadata);
+
+      LOGGER.info("Successfully registered segment: {} for table: {}", segmentName, tableNameWithType);
+      return new SuccessResponse(
+          "Successfully registered segment: " + segmentName + " of table: " + tableNameWithType);
+    } catch (WebApplicationException e) {
+      throw e;
+    } catch (Exception e) {
+      _controllerMetrics.addMeteredGlobalValue(ControllerMeter.CONTROLLER_SEGMENT_UPLOAD_ERROR, 1L);
+      _controllerMetrics.addMeteredTableValue(tableName, ControllerMeter.CONTROLLER_TABLE_SEGMENT_UPLOAD_ERROR, 1L);
+      throw new ControllerApplicationException(LOGGER,
+          "Exception while registering segment: " + e.getMessage(), Response.Status.INTERNAL_SERVER_ERROR, e);
+    }
+  }
+
   @POST
   @Path("segments/{tableName}/startReplaceSegments")
   @Authorize(targetType = TargetType.TABLE, paramName = "tableName", action = Actions.Table.REPLACE_SEGMENT)

pinot-controller/src/main/java/org/apache/pinot/controller/api/upload/ZKOperator.java

@@ -180,6 +180,64 @@ public void completeSegmentsOperations(TableConfig tableConfig, FileUploadType u
     processExistingSegments(tableConfig, uploadType, enableParallelPushProtection, headers, existingSegmentsList);
   }
 
+  /**
+   * Registers a segment directly from pre-built ZK metadata, without file upload or deep store copy.
+   * Only handles new segment creation. Throws CONFLICT if the segment already exists.
+   */
+  public void registerSegment(TableConfig tableConfig, SegmentZKMetadata segmentZKMetadata) {
+    String tableNameWithType = tableConfig.getTableName();
+    String segmentName = segmentZKMetadata.getSegmentName();
+
+    ZNRecord existingSegmentMetadataZNRecord =
+        _pinotHelixResourceManager.getSegmentMetadataZnRecord(tableNameWithType, segmentName);
+    if (existingSegmentMetadataZNRecord != null && shouldProcessAsNewSegment(tableNameWithType, segmentName,
+        existingSegmentMetadataZNRecord, false)) {
+      LOGGER.warn("Removing segment ZK metadata (recovering from previous registration failure) for table: {}, "
+          + "segment: {}", tableNameWithType, segmentName);
+      Preconditions.checkState(_pinotHelixResourceManager.removeSegmentZKMetadata(tableNameWithType, segmentName),
+          "Failed to remove segment ZK metadata for table: %s, segment: %s", tableNameWithType, segmentName);
+      existingSegmentMetadataZNRecord = null;
+    }
+
+    if (existingSegmentMetadataZNRecord != null) {
+      LOGGER.error("Segment: {} already exists in table: {}, cannot register", segmentName, tableNameWithType);
+      throw new ControllerApplicationException(LOGGER,
+          String.format("Segment: %s already exists in table: %s.", segmentName, tableNameWithType),
+          Response.Status.CONFLICT);
+    }
+
+    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, true);
+    } catch (Exception e) {
+      LOGGER.error("Failed to assign segment: {} for table: {}, ZK metadata already cleaned up by assignSegment",
+          segmentName, tableNameWithType, e);
+      throw e;
+    }
+
+    try {
+      // This call updates the sourceServer in the segment ZK metadata.
+      if (segmentZKMetadata.getSourceServer() != null) {
+        if (_pinotHelixResourceManager.updateZkMetadata(tableNameWithType, segmentZKMetadata, 0)) {
+          LOGGER.info("Set sourceServer: {} for segment: {} of table: {}",
+              segmentZKMetadata.getSourceServer(), segmentName, tableNameWithType);
+        } else {
+          LOGGER.warn(
+              "Failed to persist sourceServer for segment: {} of table: {}, servers will fall back to deep store",
+              segmentName, tableNameWithType);
+        }
+      }
+    } catch (Exception e) {
+      LOGGER.error("Failed to update sourceServer. segment {}, table: {}",
+          segmentName, tableNameWithType, e);
+    }
+  }
+
   public void completeReingestedSegmentOperations(String realtimeTableName, SegmentMetadata segmentMetadata,
       URI finalSegmentLocationURI, String sourceDownloadURIStr, String segmentDownloadURIStr, long segmentSizeInBytes)
       throws Exception {