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

Example 16 with ProcessContinuation

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

the class QueryChangeStreamActionTest method testQueryChangeStreamWithRestrictionStartAfterPartitionStart.

@Test
public void testQueryChangeStreamWithRestrictionStartAfterPartitionStart() {
    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);
    // One microsecond after partition start timestamp
    when(restriction.getFrom()).thenReturn(11L);
    // This record should be ignored because it is before restriction.getFrom
    when(record1.getRecordTimestamp()).thenReturn(Timestamp.ofTimeMicroseconds(10L));
    // This record should be included because it is at the restriction.getFrom
    when(record2.getRecordTimestamp()).thenReturn(Timestamp.ofTimeMicroseconds(11L));
    // We should start the query 1 microsecond before the restriction.getFrom
    when(changeStreamDao.changeStreamQuery(PARTITION_TOKEN, Timestamp.ofTimeMicroseconds(10L), 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, 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, record2, restrictionTracker, watermarkEstimator);
    verify(partitionMetadataDao).updateWatermark(PARTITION_TOKEN, WATERMARK_TIMESTAMP);
    verify(childPartitionsRecordAction, never()).run(partition, record1, restrictionTracker, watermarkEstimator);
    verify(dataChangeRecordAction, never()).run(any(), any(), any(), any(), any());
    verify(heartbeatRecordAction, never()).run(any(), any(), any(), any());
    verify(restrictionTracker, never()).tryClaim(any());
}

Example 17 with ProcessContinuation

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

the class QueryChangeStreamActionTest method testQueryChangeStreamWithStreamFinished.

@Test
public void testQueryChangeStreamWithStreamFinished() {
    final ChangeStreamResultSet changeStreamResultSet = mock(ChangeStreamResultSet.class);
    when(changeStreamDao.changeStreamQuery(PARTITION_TOKEN, PARTITION_START_TIMESTAMP, PARTITION_END_TIMESTAMP, PARTITION_HEARTBEAT_MILLIS)).thenReturn(changeStreamResultSet);
    when(changeStreamResultSet.next()).thenReturn(false);
    when(watermarkEstimator.currentWatermark()).thenReturn(WATERMARK);
    when(restrictionTracker.tryClaim(PARTITION_END_MICROS)).thenReturn(true);
    final ProcessContinuation result = action.run(partition, restrictionTracker, outputReceiver, watermarkEstimator, bundleFinalizer);
    assertEquals(ProcessContinuation.stop(), result);
    verify(partitionMetadataDao).updateWatermark(PARTITION_TOKEN, WATERMARK_TIMESTAMP);
    verify(partitionMetadataDao).updateToFinished(PARTITION_TOKEN);
    verify(dataChangeRecordAction, never()).run(any(), any(), any(), any(), any());
    verify(heartbeatRecordAction, never()).run(any(), any(), any(), any());
    verify(childPartitionsRecordAction, never()).run(any(), any(), any(), any());
}

Example 18 with ProcessContinuation

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

the class ChildPartitionsRecordActionTest method testRestrictionClaimedAndIsMergeCaseAndChildExists.

@Test
public void testRestrictionClaimedAndIsMergeCaseAndChildExists() {
    final String partitionToken = "partitionToken";
    final String anotherPartitionToken = "anotherPartitionToken";
    final String childPartitionToken = "childPartition1";
    final HashSet<String> parentTokens = Sets.newHashSet(partitionToken, anotherPartitionToken);
    final long heartbeat = 30L;
    final Timestamp startTimestamp = Timestamp.ofTimeMicroseconds(10L);
    final Timestamp endTimestamp = Timestamp.ofTimeMicroseconds(20L);
    final PartitionMetadata partition = mock(PartitionMetadata.class);
    final ChildPartitionsRecord record = new ChildPartitionsRecord(startTimestamp, "recordSequence", Collections.singletonList(new ChildPartition(childPartitionToken, parentTokens)), null);
    when(partition.getEndTimestamp()).thenReturn(endTimestamp);
    when(partition.getHeartbeatMillis()).thenReturn(heartbeat);
    when(partition.getPartitionToken()).thenReturn(partitionToken);
    when(tracker.tryClaim(10L)).thenReturn(true);
    when(transaction.getPartition(childPartitionToken)).thenReturn(mock(Struct.class));
    final Optional<ProcessContinuation> maybeContinuation = action.run(partition, record, tracker, watermarkEstimator);
    assertEquals(Optional.empty(), maybeContinuation);
    verify(watermarkEstimator).setWatermark(new Instant(startTimestamp.toSqlTimestamp().getTime()));
    verify(transaction, never()).insert(any());
}

Example 19 with ProcessContinuation

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

the class ChildPartitionsRecordActionTest method testRestrictionClaimedAndIsSplitCase.

@Test
public void testRestrictionClaimedAndIsSplitCase() {
    final String partitionToken = "partitionToken";
    final long heartbeat = 30L;
    final Timestamp startTimestamp = Timestamp.ofTimeMicroseconds(10L);
    final Timestamp endTimestamp = Timestamp.ofTimeMicroseconds(20L);
    final PartitionMetadata partition = mock(PartitionMetadata.class);
    final ChildPartitionsRecord record = new ChildPartitionsRecord(startTimestamp, "recordSequence", Arrays.asList(new ChildPartition("childPartition1", partitionToken), new ChildPartition("childPartition2", partitionToken)), null);
    when(partition.getEndTimestamp()).thenReturn(endTimestamp);
    when(partition.getHeartbeatMillis()).thenReturn(heartbeat);
    when(partition.getPartitionToken()).thenReturn(partitionToken);
    when(tracker.tryClaim(10L)).thenReturn(true);
    when(transaction.getPartition("childPartition1")).thenReturn(null);
    when(transaction.getPartition("childPartition2")).thenReturn(null);
    final Optional<ProcessContinuation> maybeContinuation = action.run(partition, record, tracker, watermarkEstimator);
    assertEquals(Optional.empty(), maybeContinuation);
    verify(watermarkEstimator).setWatermark(new Instant(startTimestamp.toSqlTimestamp().getTime()));
    verify(transaction).insert(PartitionMetadata.newBuilder().setPartitionToken("childPartition1").setParentTokens(Sets.newHashSet(partitionToken)).setStartTimestamp(startTimestamp).setEndTimestamp(endTimestamp).setHeartbeatMillis(heartbeat).setState(CREATED).setWatermark(startTimestamp).build());
    verify(transaction).insert(PartitionMetadata.newBuilder().setPartitionToken("childPartition2").setParentTokens(Sets.newHashSet(partitionToken)).setStartTimestamp(startTimestamp).setEndTimestamp(endTimestamp).setHeartbeatMillis(heartbeat).setState(CREATED).setWatermark(startTimestamp).build());
}

Example 20 with ProcessContinuation

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

the class ChildPartitionsRecordActionTest method testRestrictionClaimedAndIsMergeCaseAndChildNotExists.

@Test
public void testRestrictionClaimedAndIsMergeCaseAndChildNotExists() {
    final String partitionToken = "partitionToken";
    final String anotherPartitionToken = "anotherPartitionToken";
    final String childPartitionToken = "childPartition1";
    final HashSet<String> parentTokens = Sets.newHashSet(partitionToken, anotherPartitionToken);
    final long heartbeat = 30L;
    final Timestamp startTimestamp = Timestamp.ofTimeMicroseconds(10L);
    final Timestamp endTimestamp = Timestamp.ofTimeMicroseconds(20L);
    final PartitionMetadata partition = mock(PartitionMetadata.class);
    final ChildPartitionsRecord record = new ChildPartitionsRecord(startTimestamp, "recordSequence", Collections.singletonList(new ChildPartition(childPartitionToken, parentTokens)), null);
    when(partition.getEndTimestamp()).thenReturn(endTimestamp);
    when(partition.getHeartbeatMillis()).thenReturn(heartbeat);
    when(partition.getPartitionToken()).thenReturn(partitionToken);
    when(tracker.tryClaim(10L)).thenReturn(true);
    when(transaction.getPartition(childPartitionToken)).thenReturn(null);
    final Optional<ProcessContinuation> maybeContinuation = action.run(partition, record, tracker, watermarkEstimator);
    assertEquals(Optional.empty(), maybeContinuation);
    verify(watermarkEstimator).setWatermark(new Instant(startTimestamp.toSqlTimestamp().getTime()));
    verify(transaction).insert(PartitionMetadata.newBuilder().setPartitionToken(childPartitionToken).setParentTokens(parentTokens).setStartTimestamp(startTimestamp).setEndTimestamp(endTimestamp).setHeartbeatMillis(heartbeat).setState(CREATED).setWatermark(startTimestamp).build());
}