public void completeSegmentOperations(String tableNameWithType, SegmentMetadata segmentMetadata,
      FileUploadType uploadType, @Nullable URI finalSegmentLocationURI, File segmentFile,
      @Nullable String sourceDownloadURIStr, String segmentDownloadURIStr, @Nullable String crypterName,
      long segmentSizeInBytes, boolean enableParallelPushProtection, boolean allowRefresh, HttpHeaders headers)
      throws Exception {
    String segmentName = segmentMetadata.getName();
    boolean refreshOnly =
        Boolean.parseBoolean(headers.getHeaderString(FileUploadDownloadClient.CustomHeaders.REFRESH_ONLY));

    ZNRecord existingSegmentMetadataZNRecord =
        _pinotHelixResourceManager.getSegmentMetadataZnRecord(tableNameWithType, segmentName);
    if (existingSegmentMetadataZNRecord != null && shouldProcessAsNewSegment(tableNameWithType, segmentName,
        existingSegmentMetadataZNRecord, enableParallelPushProtection)) {
      LOGGER.warn("Removing segment ZK metadata (recovering from previous upload 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) {
      // Add a new segment
      if (refreshOnly) {
        throw new ControllerApplicationException(LOGGER,
            String.format("Cannot refresh non-existing segment: %s for table: %s", segmentName, tableNameWithType),
            Response.Status.GONE);
      }
      LOGGER.info("Adding new segment: {} to table: {}", segmentName, tableNameWithType);
      processNewSegment(tableNameWithType, segmentMetadata, uploadType, finalSegmentLocationURI, segmentFile,
          sourceDownloadURIStr, segmentDownloadURIStr, crypterName, segmentSizeInBytes, enableParallelPushProtection,
          headers);
    } else {
      // Refresh an existing segment
      if (!allowRefresh) {
        // We cannot perform this check up-front in UploadSegment API call. If a segment doesn't exist during the check
        // done up-front but ends up getting created before the check here, we could incorrectly refresh an existing
        // segment.
        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 it", segmentName, tableNameWithType);
      processExistingSegment(tableNameWithType, segmentMetadata, uploadType, existingSegmentMetadataZNRecord,
          finalSegmentLocationURI, segmentFile, sourceDownloadURIStr, segmentDownloadURIStr, crypterName,
          segmentSizeInBytes, enableParallelPushProtection, headers);
    }
  }
@Test
  public void testMetadataUploadType()
      throws Exception {
    String segmentName = "metadataTest";
    FileUtils.deleteQuietly(TEMP_DIR);
    ZKOperator zkOperator = new ZKOperator(_resourceManager, mock(ControllerConf.class), mock(ControllerMetrics.class));

    SegmentMetadata segmentMetadata = mock(SegmentMetadata.class);
    when(segmentMetadata.getName()).thenReturn(segmentName);
    when(segmentMetadata.getCrc()).thenReturn("12345");
    when(segmentMetadata.getIndexCreationTime()).thenReturn(123L);
    HttpHeaders httpHeaders = mock(HttpHeaders.class);

    File segmentTar = generateSegment();
    String sourceDownloadURIStr = segmentTar.toURI().toString();
    File segmentFile = new File("metadataOnly");

    // with finalSegmentLocation not null
    File finalSegmentLocation = new File(DATA_DIR, segmentName);
    Assert.assertFalse(finalSegmentLocation.exists());
    zkOperator.completeSegmentOperations(OFFLINE_TABLE_NAME, segmentMetadata, FileUploadType.METADATA,
        finalSegmentLocation.toURI(), segmentFile, sourceDownloadURIStr, "downloadUrl", "crypter", 10, true, true,
        httpHeaders);
    Assert.assertTrue(finalSegmentLocation.exists());
    Assert.assertTrue(segmentTar.exists());
    checkSegmentZkMetadata(segmentName, 12345L, 123L);

    _resourceManager.deleteSegment(OFFLINE_TABLE_NAME, segmentName);
    // Wait for the segment Zk entry to be deleted.
    TestUtils.waitForCondition(aVoid -> {
      SegmentZKMetadata segmentZKMetadata = _resourceManager.getSegmentZKMetadata(OFFLINE_TABLE_NAME, segmentName);
      return segmentZKMetadata == null;
    }, 30_000L, "Failed to delete segmentZkMetadata.");

    FileUtils.deleteQuietly(DATA_DIR);
    // with finalSegmentLocation null
    zkOperator.completeSegmentOperations(OFFLINE_TABLE_NAME, segmentMetadata, FileUploadType.METADATA, null,
        segmentFile, sourceDownloadURIStr, "downloadUrl", "crypter", 10, true, true, httpHeaders);
    Assert.assertFalse(finalSegmentLocation.exists());
    Assert.assertTrue(segmentTar.exists());
    checkSegmentZkMetadata(segmentName, 12345L, 123L);
  }