diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java index af77b6ab580d..f8177af78eaa 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java @@ -367,16 +367,8 @@ private SuccessResponse uploadSegment(@Nullable String tableName, TableType tabl // Fetch segment name String segmentName = segmentMetadata.getName(); - // Fetch table name. Try to derive the table name from the parameter and then from segment metadata - String rawTableName; - if (StringUtils.isNotEmpty(tableName)) { - rawTableName = TableNameBuilder.extractRawTableName(tableName); - } else { - // TODO: remove this when we completely deprecate the table name from segment metadata - rawTableName = segmentMetadata.getTableName(); - LOGGER.warn("Table name is not provided as request query parameter when uploading segment: {} for table: {}", - segmentName, rawTableName); - } + // Prefer request tableName over table name baked into segment metadata (enables staging→prod promote) + String rawTableName = resolveRawTableNameForUpload(tableName, segmentMetadata); String tableNameWithType = tableType == TableType.OFFLINE ? TableNameBuilder.OFFLINE.tableNameWithType(rawTableName) : TableNameBuilder.REALTIME.tableNameWithType(rawTableName); @@ -830,7 +822,9 @@ private void decryptFile(String crypterClassName, File tempEncryptedFile, File t // request if a multipart object is not sent. This endpoint does not move the segment to its final location; // it keeps it at the downloadURI header that is set. We will not support this endpoint going forward. public void uploadSegmentAsJson(String segmentJsonStr, - @ApiParam(value = "Name of the table") @QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_NAME) + @ApiParam(value = "Name of the table to upload into. Overrides segment.table.name in segment metadata when set " + + "(allows promoting a segment built for another table). Falls back to metadata when omitted.") + @QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_NAME) String tableName, @ApiParam(value = "Type of the table") @QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_TYPE) @DefaultValue("OFFLINE") String tableType, @@ -869,7 +863,9 @@ public void uploadSegmentAsJson(String segmentJsonStr, @TrackedByGauge(gauge = ControllerGauge.SEGMENT_UPLOADS_IN_PROGRESS) // For the multipart endpoint, we will always move segment to final location regardless of the segment endpoint. public void uploadSegmentAsMultiPart(FormDataMultiPart multiPart, - @ApiParam(value = "Name of the table") @QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_NAME) + @ApiParam(value = "Name of the table to upload into. Overrides segment.table.name in segment metadata when set " + + "(allows promoting a segment built for another table). Falls back to metadata when omitted.") + @QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_NAME) String tableName, @ApiParam(value = "Type of the table") @QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_TYPE) @DefaultValue("OFFLINE") String tableType, @@ -971,7 +967,8 @@ public void uploadSegmentsAsMultiPart(FormDataMultiPart multiPart, // request if a multipart object is not sent. This endpoint is recommended for use. It differs from the first // endpoint in how it moves the segment to a Pinot-determined final directory. public void uploadSegmentAsJsonV2(String segmentJsonStr, - @ApiParam(value = "Name of the table") @QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_NAME) + @ApiParam(value = "Name of the table to upload into. Overrides segment.table.name in segment metadata when set.") + @QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_NAME) String tableName, @ApiParam(value = "Type of the table") @QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_TYPE) @DefaultValue("OFFLINE") String tableType, @@ -1011,7 +1008,8 @@ public void uploadSegmentAsJsonV2(String segmentJsonStr, @TrackedByGauge(gauge = ControllerGauge.SEGMENT_UPLOADS_IN_PROGRESS) // This behavior does not differ from v1 of the same endpoint. public void uploadSegmentAsMultiPartV2(FormDataMultiPart multiPart, - @ApiParam(value = "Name of the table") @QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_NAME) + @ApiParam(value = "Name of the table to upload into. Overrides segment.table.name in segment metadata when set.") + @QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_NAME) String tableName, @ApiParam(value = "Type of the table") @QueryParam(FileUploadDownloadClient.QueryParameters.TABLE_TYPE) @DefaultValue("OFFLINE") String tableType, @@ -1313,6 +1311,38 @@ private FileUploadType getUploadType(String uploadTypeStr) { } } + /** + * Resolve the raw table name for a segment upload. + *

+ * The request {@code tableName} query parameter (also set by admin {@code -tableName} and batch + * {@code tableSpec.tableName}) is authoritative. Segment metadata {@code segment.table.name} is only + * used when the request omits the parameter. Callers may therefore build a segment for table A and + * upload it to table B without rewriting the tar. + * + * @param requestTableName table name from the upload request (may be null/empty or typed) + * @param segmentMetadata metadata read from the uploaded segment + * @return raw (untyped) table name used for ZK / deep-store routing + */ + @VisibleForTesting + static String resolveRawTableNameForUpload(@Nullable String requestTableName, SegmentMetadata segmentMetadata) { + String metadataTableName = segmentMetadata.getTableName(); + if (StringUtils.isNotEmpty(requestTableName)) { + String rawTableName = TableNameBuilder.extractRawTableName(requestTableName); + if (StringUtils.isNotEmpty(metadataTableName)) { + String metadataRawTableName = TableNameBuilder.extractRawTableName(metadataTableName); + if (!rawTableName.equals(metadataRawTableName)) { + LOGGER.info("Uploading segment: {} with request table name: {} which differs from segment metadata " + + "table name: {}", segmentMetadata.getName(), rawTableName, metadataRawTableName); + } + } + return rawTableName; + } + // TODO: remove this when we completely deprecate the table name from segment metadata + LOGGER.warn("Table name is not provided as request query parameter when uploading segment: {} for table: {}", + segmentMetadata.getName(), metadataTableName); + return metadataTableName; + } + @VisibleForTesting static long getSegmentSizeFromFile(String sourceDownloadURIStr) throws IOException { diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResourceTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResourceTest.java index 60aaa346d46a..02a01d3eb182 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResourceTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResourceTest.java @@ -36,6 +36,8 @@ import org.apache.pinot.controller.ControllerConf; import org.apache.pinot.controller.api.exception.ControllerApplicationException; import org.apache.pinot.controller.api.upload.SegmentMetadataInfo; +import org.apache.pinot.controller.utils.SegmentMetadataMockUtils; +import org.apache.pinot.segment.spi.SegmentMetadata; import org.apache.pinot.spi.crypt.NoOpPinotCrypter; import org.apache.pinot.spi.crypt.PinotCrypterFactory; import org.apache.pinot.spi.env.PinotConfiguration; @@ -313,4 +315,21 @@ public void testCreateSegmentFileFromMultipart() throw new AssertionError("Method threw an exception: " + e.getMessage(), e); } } + + /** + * Request tableName is authoritative over segment.table.name baked into metadata. + * Enables build-for-A / upload-to-B (staging → prod promote) without rewriting the tar. + */ + @Test + public void testResolveRawTableNamePrefersRequestOverMetadata() { + SegmentMetadata metadata = SegmentMetadataMockUtils.mockSegmentMetadata("tableA", "seg1"); + + assertEquals(PinotSegmentUploadDownloadRestletResource.resolveRawTableNameForUpload("tableB", metadata), + "tableB"); + assertEquals(PinotSegmentUploadDownloadRestletResource.resolveRawTableNameForUpload("tableB_OFFLINE", metadata), + "tableB"); + assertEquals(PinotSegmentUploadDownloadRestletResource.resolveRawTableNameForUpload(null, metadata), "tableA"); + assertEquals(PinotSegmentUploadDownloadRestletResource.resolveRawTableNameForUpload("", metadata), "tableA"); + assertEquals(PinotSegmentUploadDownloadRestletResource.resolveRawTableNameForUpload("tableA", metadata), "tableA"); + } } diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/UploadSegmentCommand.java b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/UploadSegmentCommand.java index 5a35d8862219..25aea372c437 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/UploadSegmentCommand.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/UploadSegmentCommand.java @@ -75,7 +75,9 @@ public class UploadSegmentCommand extends AbstractBaseAdminCommand implements Co @CommandLine.Option(names = {"-segmentDir"}, required = true, description = "Path to segment directory.") private String _segmentDir = null; - @CommandLine.Option(names = {"-tableName"}, required = false, description = "Table name to upload") + @CommandLine.Option(names = {"-tableName"}, required = false, + description = "Table name to upload into. Overrides segment.table.name baked into the segment metadata " + + "(build once, push to staging or prod). When omitted, falls back to the name in segment metadata.") private String _tableName = null; @CommandLine.Option(names = {"-tableType"}, required = false,