Examples with ProcessContinuation org.apache.beam.sdk.transforms.DoFn.ProcessContinuation used on opensource projects

Example 11 with ProcessContinuation

use of org.apache.beam.sdk.transforms.DoFn.ProcessContinuation in project beam by apache.

the class DataChangeRecordActionTest method testRestrictionNotClaimed.

@Test
public void testRestrictionNotClaimed() {
    final String partitionToken = "partitionToken";
    final Timestamp timestamp = Timestamp.ofTimeMicroseconds(10L);
    final DataChangeRecord record = mock(DataChangeRecord.class);
    when(record.getCommitTimestamp()).thenReturn(timestamp);
    when(tracker.tryClaim(10L)).thenReturn(false);
    when(partition.getPartitionToken()).thenReturn(partitionToken);
    final Optional<ProcessContinuation> maybeContinuation = action.run(partition, record, tracker, outputReceiver, watermarkEstimator);
    assertEquals(Optional.of(ProcessContinuation.stop()), maybeContinuation);
    verify(outputReceiver, never()).outputWithTimestamp(any(), any());
    verify(watermarkEstimator, never()).setWatermark(any());
}

Example 12 with ProcessContinuation

use of org.apache.beam.sdk.transforms.DoFn.ProcessContinuation in project beam by apache.

the class QueryChangeStreamActionTest method testQueryChangeStreamWithDataChangeRecord.

@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, rowAsStruct, 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());
}

Example 13 with ProcessContinuation

use of org.apache.beam.sdk.transforms.DoFn.ProcessContinuation in project beam by apache.

the class QueryChangeStreamActionTest method testQueryChangeStreamWithHeartbeatRecord.

@Test
public void testQueryChangeStreamWithHeartbeatRecord() {
    final Struct rowAsStruct = mock(Struct.class);
    final ChangeStreamResultSetMetadata resultSetMetadata = mock(ChangeStreamResultSetMetadata.class);
    final ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class);
    final HeartbeatRecord record1 = mock(HeartbeatRecord.class);
    final HeartbeatRecord record2 = mock(HeartbeatRecord.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, rowAsStruct, resultSetMetadata)).thenReturn(Arrays.asList(record1, record2));
    when(heartbeatRecordAction.run(partition, record1, restrictionTracker, watermarkEstimator)).thenReturn(Optional.empty());
    when(heartbeatRecordAction.run(partition, record2, restrictionTracker, 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(heartbeatRecordAction).run(partition, record1, restrictionTracker, watermarkEstimator);
    verify(heartbeatRecordAction).run(partition, record2, restrictionTracker, watermarkEstimator);
    verify(partitionMetadataDao).updateWatermark(PARTITION_TOKEN, WATERMARK_TIMESTAMP);
    verify(dataChangeRecordAction, never()).run(any(), any(), any(), any(), any());
    verify(childPartitionsRecordAction, never()).run(any(), any(), any(), any());
    verify(restrictionTracker, never()).tryClaim(any());
}

Example 14 with ProcessContinuation

use of org.apache.beam.sdk.transforms.DoFn.ProcessContinuation in project beam by apache.

the class DataChangeRecordActionTest method testRestrictionClaimed.

@Test
public void testRestrictionClaimed() {
    final String partitionToken = "partitionToken";
    final Timestamp timestamp = Timestamp.ofTimeMicroseconds(10L);
    final Instant instant = new Instant(timestamp.toSqlTimestamp().getTime());
    final DataChangeRecord record = mock(DataChangeRecord.class);
    when(record.getCommitTimestamp()).thenReturn(timestamp);
    when(tracker.tryClaim(10L)).thenReturn(true);
    when(partition.getPartitionToken()).thenReturn(partitionToken);
    final Optional<ProcessContinuation> maybeContinuation = action.run(partition, record, tracker, outputReceiver, watermarkEstimator);
    assertEquals(Optional.empty(), maybeContinuation);
    verify(outputReceiver).outputWithTimestamp(record, instant);
    verify(watermarkEstimator).setWatermark(instant);
}

Example 15 with ProcessContinuation

use of org.apache.beam.sdk.transforms.DoFn.ProcessContinuation in project beam by apache.

the class QueryChangeStreamActionTest method testQueryChangeStreamWithChildPartitionsRecord.

@Test
public void testQueryChangeStreamWithChildPartitionsRecord() {
    final Struct rowAsStruct = mock(Struct.class);
    final ChangeStreamResultSetMetadata resultSetMetadata = mock(ChangeStreamResultSetMetadata.class);
    final ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class);
    final ChildPartitionsRecord record1 = mock(ChildPartitionsRecord.class);
    final ChildPartitionsRecord record2 = mock(ChildPartitionsRecord.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, rowAsStruct, resultSetMetadata)).thenReturn(Arrays.asList(record1, record2));
    when(childPartitionsRecordAction.run(partition, record1, restrictionTracker, watermarkEstimator)).thenReturn(Optional.empty());
    when(childPartitionsRecordAction.run(partition, record2, restrictionTracker, 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(childPartitionsRecordAction).run(partition, record1, restrictionTracker, watermarkEstimator);
    verify(childPartitionsRecordAction).run(partition, record2, restrictionTracker, watermarkEstimator);
    verify(partitionMetadataDao).updateWatermark(PARTITION_TOKEN, WATERMARK_TIMESTAMP);
    verify(dataChangeRecordAction, never()).run(any(), any(), any(), any(), any());
    verify(heartbeatRecordAction, never()).run(any(), any(), any(), any());
    verify(restrictionTracker, never()).tryClaim(any());
}