Examples with ChildPartitionsRecord org.apache.beam.sdk.io.gcp.spanner.changestreams.model.ChildPartitionsRecord used on opensource projects

Example 6 with ChildPartitionsRecord

use of org.apache.beam.sdk.io.gcp.spanner.changestreams.model.ChildPartitionsRecord 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 7 with ChildPartitionsRecord

use of org.apache.beam.sdk.io.gcp.spanner.changestreams.model.ChildPartitionsRecord in project beam by apache.

the class ChangeStreamRecordMapperTest method testMappingStructRowFromInitialPartitionToChildPartitionRecord.

/**
 * Adds the default parent partition token as a parent of each child partition.
 */
@Test
public void testMappingStructRowFromInitialPartitionToChildPartitionRecord() {
    final Struct struct = recordsToStructWithStrings(new ChildPartitionsRecord(Timestamp.ofTimeSecondsAndNanos(10L, 20), "1", Arrays.asList(new ChildPartition("childToken1", Sets.newHashSet()), new ChildPartition("childToken2", Sets.newHashSet())), null));
    final ChildPartitionsRecord expected = new ChildPartitionsRecord(Timestamp.ofTimeSecondsAndNanos(10L, 20), "1", Arrays.asList(new ChildPartition("childToken1", Sets.newHashSet(InitialPartition.PARTITION_TOKEN)), new ChildPartition("childToken2", Sets.newHashSet(InitialPartition.PARTITION_TOKEN))), null);
    final PartitionMetadata initialPartition = partition.toBuilder().setPartitionToken(InitialPartition.PARTITION_TOKEN).build();
    assertEquals(Collections.singletonList(expected), mapper.toChangeStreamRecords(initialPartition, struct, resultSetMetadata));
}

Example 8 with ChildPartitionsRecord

use of org.apache.beam.sdk.io.gcp.spanner.changestreams.model.ChildPartitionsRecord in project beam by apache.

the class ChangeStreamRecordMapperTest method testMappingStructRowToChildPartitionRecord.

@Test
public void testMappingStructRowToChildPartitionRecord() {
    final ChildPartitionsRecord childPartitionsRecord = new ChildPartitionsRecord(Timestamp.ofTimeSecondsAndNanos(10L, 20), "1", Arrays.asList(new ChildPartition("childToken1", Sets.newHashSet("parentToken1", "parentToken2")), new ChildPartition("childToken2", Sets.newHashSet("parentToken1", "parentToken2"))), null);
    final Struct struct = recordsToStructWithStrings(childPartitionsRecord);
    assertEquals(Collections.singletonList(childPartitionsRecord), mapper.toChangeStreamRecords(partition, struct, resultSetMetadata));
}

Example 9 with ChildPartitionsRecord

use of org.apache.beam.sdk.io.gcp.spanner.changestreams.model.ChildPartitionsRecord 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 10 with ChildPartitionsRecord

use of org.apache.beam.sdk.io.gcp.spanner.changestreams.model.ChildPartitionsRecord 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());
}