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 testPushToRealtimeTable()
      throws Exception {
    ZKOperator zkOperator = new ZKOperator(_resourceManager, mock(ControllerConf.class), mock(ControllerMetrics.class));

    SegmentMetadata segmentMetadata = mock(SegmentMetadata.class);
    when(segmentMetadata.getName()).thenReturn(SEGMENT_NAME);
    when(segmentMetadata.getCrc()).thenReturn("12345");
    zkOperator.completeSegmentOperations(REALTIME_TABLE_NAME, segmentMetadata, FileUploadType.SEGMENT, null, null,
        "downloadUrl", "downloadUrl", null, 10, true, true, mock(HttpHeaders.class));

    SegmentZKMetadata segmentZKMetadata = _resourceManager.getSegmentZKMetadata(REALTIME_TABLE_NAME, SEGMENT_NAME);
    assertNotNull(segmentZKMetadata);
    assertEquals(segmentZKMetadata.getStatus(), Status.UPLOADED);
    assertNull(segmentMetadata.getStartOffset());
    assertNull(segmentMetadata.getEndOffset());

    // Uploading a segment with LLC segment name but without start/end offset should fail
    when(segmentMetadata.getName()).thenReturn(LLC_SEGMENT_NAME);
    when(segmentMetadata.getCrc()).thenReturn("23456");
    try {
      zkOperator.completeSegmentOperations(REALTIME_TABLE_NAME, segmentMetadata, FileUploadType.SEGMENT, null, null,
          "downloadUrl", "downloadUrl", null, 10, true, true, mock(HttpHeaders.class));
      fail();
    } catch (ControllerApplicationException e) {
      assertEquals(e.getResponse().getStatus(), Response.Status.BAD_REQUEST.getStatusCode());
    }
    assertNull(_resourceManager.getSegmentZKMetadata(REALTIME_TABLE_NAME, LLC_SEGMENT_NAME));

    // Uploading a segment with LLC segment name and start/end offset should success
    when(segmentMetadata.getStartOffset()).thenReturn("0");
    when(segmentMetadata.getEndOffset()).thenReturn("1234");
    zkOperator.completeSegmentOperations(REALTIME_TABLE_NAME, segmentMetadata, FileUploadType.SEGMENT, null, null,
        "downloadUrl", "downloadUrl", null, 10, true, true, mock(HttpHeaders.class));

    segmentZKMetadata = _resourceManager.getSegmentZKMetadata(REALTIME_TABLE_NAME, LLC_SEGMENT_NAME);
    assertNotNull(segmentZKMetadata);
    assertEquals(segmentZKMetadata.getStatus(), Status.UPLOADED);
    assertEquals(segmentZKMetadata.getStartOffset(), "0");
    assertEquals(segmentZKMetadata.getEndOffset(), "1234");

    // Refreshing a segment with LLC segment name but without start/end offset should success
    when(segmentMetadata.getCrc()).thenReturn("34567");
    when(segmentMetadata.getStartOffset()).thenReturn(null);
    when(segmentMetadata.getEndOffset()).thenReturn(null);
    zkOperator.completeSegmentOperations(REALTIME_TABLE_NAME, segmentMetadata, FileUploadType.SEGMENT, null, null,
        "downloadUrl", "downloadUrl", null, 10, true, true, mock(HttpHeaders.class));

    segmentZKMetadata = _resourceManager.getSegmentZKMetadata(REALTIME_TABLE_NAME, LLC_SEGMENT_NAME);
    assertNotNull(segmentZKMetadata);
    assertEquals(segmentZKMetadata.getStatus(), Status.UPLOADED);
    assertEquals(segmentZKMetadata.getStartOffset(), "0");
    assertEquals(segmentZKMetadata.getEndOffset(), "1234");

    // Refreshing a segment with LLC segment name and start/end offset should override the offsets
    when(segmentMetadata.getCrc()).thenReturn("45678");
    when(segmentMetadata.getStartOffset()).thenReturn("1234");
    when(segmentMetadata.getEndOffset()).thenReturn("2345");
    zkOperator.completeSegmentOperations(REALTIME_TABLE_NAME, segmentMetadata, FileUploadType.SEGMENT, null, null,
        "downloadUrl", "downloadUrl", null, 10, true, true, mock(HttpHeaders.class));

    segmentZKMetadata = _resourceManager.getSegmentZKMetadata(REALTIME_TABLE_NAME, LLC_SEGMENT_NAME);
    assertNotNull(segmentZKMetadata);
    assertEquals(segmentZKMetadata.getStatus(), Status.UPLOADED);
    assertEquals(segmentZKMetadata.getStartOffset(), "1234");
    assertEquals(segmentZKMetadata.getEndOffset(), "2345");
  }