@SuppressWarnings("nullness")
  @VisibleForTesting
  public ProcessContinuation run(
      PartitionMetadata partition,
      RestrictionTracker<TimestampRange, Timestamp> tracker,
      OutputReceiver<DataChangeRecord> receiver,
      ManualWatermarkEstimator<Instant> watermarkEstimator,
      BundleFinalizer bundleFinalizer) {
    final String token = partition.getPartitionToken();
    final Timestamp startTimestamp = tracker.currentRestriction().getFrom();
    final Timestamp endTimestamp = partition.getEndTimestamp();

    // TODO: Potentially we can avoid this fetch, by enriching the runningAt timestamp when the
    // ReadChangeStreamPartitionDoFn#processElement is called
    final PartitionMetadata updatedPartition =
        Optional.ofNullable(partitionMetadataDao.getPartition(token))
            .map(partitionMetadataMapper::from)
            .orElseThrow(
                () ->
                    new IllegalStateException(
                        "Partition " + token + " not found in metadata table"));

    try (ChangeStreamResultSet resultSet =
        changeStreamDao.changeStreamQuery(
            token, startTimestamp, endTimestamp, partition.getHeartbeatMillis())) {

      metrics.incQueryCounter();
      while (resultSet.next()) {
        final List<ChangeStreamRecord> records =
            changeStreamRecordMapper.toChangeStreamRecords(
                updatedPartition, resultSet, resultSet.getMetadata());

        Optional<ProcessContinuation> maybeContinuation;
        for (final ChangeStreamRecord record : records) {
          if (record instanceof DataChangeRecord) {
            maybeContinuation =
                dataChangeRecordAction.run(
                    updatedPartition,
                    (DataChangeRecord) record,
                    tracker,
                    receiver,
                    watermarkEstimator);
          } else if (record instanceof HeartbeatRecord) {
            maybeContinuation =
                heartbeatRecordAction.run(
                    updatedPartition, (HeartbeatRecord) record, tracker, watermarkEstimator);
          } else if (record instanceof ChildPartitionsRecord) {
            maybeContinuation =
                childPartitionsRecordAction.run(
                    updatedPartition, (ChildPartitionsRecord) record, tracker, watermarkEstimator);
          } else {
            LOG.error("[{}] Unknown record type {}", token, record.getClass());
            throw new IllegalArgumentException("Unknown record type " + record.getClass());
          }

          if (maybeContinuation.isPresent()) {
            LOG.debug("[{}] Continuation present, returning {}", token, maybeContinuation);
            bundleFinalizer.afterBundleCommit(
                Instant.now().plus(BUNDLE_FINALIZER_TIMEOUT),
                updateWatermarkCallback(token, watermarkEstimator));
            return maybeContinuation.get();
          }
        }
      }
      bundleFinalizer.afterBundleCommit(
          Instant.now().plus(BUNDLE_FINALIZER_TIMEOUT),
          updateWatermarkCallback(token, watermarkEstimator));

    } catch (SpannerException e) {
      /*
      If there is a split when a partition is supposed to be finished, the residual will try
      to perform a change stream query for an out of range interval. We ignore this error
      here, and the residual should be able to claim the end of the timestamp range, finishing
      the partition.
      */
      if (isTimestampOutOfRange(e)) {
        LOG.info(
            "[{}] query change stream is out of range for {} to {}, finishing stream.",
            token,
            startTimestamp,
            endTimestamp,
            e);
      } else {
        throw e;
      }
    } catch (Exception e) {
      LOG.error(
          "[{}] query change stream had exception processing range {} to {}.",
          token,
          startTimestamp,
          endTimestamp,
          e);
      throw e;
    }

    LOG.debug("[{}] change stream completed successfully", token);
    if (tracker.tryClaim(endTimestamp)) {
      LOG.debug("[{}] Finishing partition", token);
      partitionMetadataDao.updateToFinished(token);
      metrics.decActivePartitionReadCounter();
      LOG.info("[{}] After attempting to finish the partition", token);
    }
    return ProcessContinuation.stop();
  }
@Test
  public void testQueryChangeStreamWithDataChangeRecord() {
    final Struct rowAsStruct = mock(Struct.class);
    final ChangeStreamResultSetMetadata resultSetMetadata =
        mock(ChangeStreamResultSetMetadata.class);
    final ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class);
    final DataChangeRecord record1 = mock(DataChangeRecord.class);
    final DataChangeRecord record2 = mock(DataChangeRecord.class);
    when(record1.getRecordTimestamp()).thenReturn(PARTITION_START_TIMESTAMP);
    when(record2.getRecordTimestamp()).thenReturn(PARTITION_START_TIMESTAMP);
    when(changeStreamDao.changeStreamQuery(
            PARTITION_TOKEN,
            PARTITION_START_TIMESTAMP,
            PARTITION_END_TIMESTAMP,
            PARTITION_HEARTBEAT_MILLIS))
        .thenReturn(resultSet);
    when(resultSet.next()).thenReturn(true);
    when(resultSet.getCurrentRowAsStruct()).thenReturn(rowAsStruct);
    when(resultSet.getMetadata()).thenReturn(resultSetMetadata);
    when(changeStreamRecordMapper.toChangeStreamRecords(partition, resultSet, resultSetMetadata))
        .thenReturn(Arrays.asList(record1, record2));
    when(dataChangeRecordAction.run(
            partition, record1, restrictionTracker, outputReceiver, watermarkEstimator))
        .thenReturn(Optional.empty());
    when(dataChangeRecordAction.run(
            partition, record2, restrictionTracker, outputReceiver, watermarkEstimator))
        .thenReturn(Optional.of(ProcessContinuation.stop()));
    when(watermarkEstimator.currentWatermark()).thenReturn(WATERMARK);

    final ProcessContinuation result =
        action.run(
            partition, restrictionTracker, outputReceiver, watermarkEstimator, bundleFinalizer);

    assertEquals(ProcessContinuation.stop(), result);
    verify(dataChangeRecordAction)
        .run(partition, record1, restrictionTracker, outputReceiver, watermarkEstimator);
    verify(dataChangeRecordAction)
        .run(partition, record2, restrictionTracker, outputReceiver, watermarkEstimator);
    verify(partitionMetadataDao).updateWatermark(PARTITION_TOKEN, WATERMARK_TIMESTAMP);

    verify(heartbeatRecordAction, never()).run(any(), any(), any(), any());
    verify(childPartitionsRecordAction, never()).run(any(), any(), any(), any());
    verify(restrictionTracker, never()).tryClaim(any());
  }