Document and test table name override on segment push

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.
+   * <p>
+   * 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 {

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");
+  }
 }

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,