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

Example 6 with PartitionMetadata

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

the class ChildPartitionsRecordActionTest method testRestrictionClaimedAnsIsSplitCaseAndChildExists.

@Test
public void testRestrictionClaimedAnsIsSplitCaseAndChildExists() {
    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(mock(Struct.class));
    when(transaction.getPartition("childPartition2")).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()));
}

Example 7 with PartitionMetadata

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

the class ChildPartitionsRecordActionTest method testRestrictionNotClaimed.

@Test
public void testRestrictionNotClaimed() {
    final String partitionToken = "partitionToken";
    final Timestamp startTimestamp = Timestamp.ofTimeMicroseconds(10L);
    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.getPartitionToken()).thenReturn(partitionToken);
    when(tracker.tryClaim(10L)).thenReturn(false);
    final Optional<ProcessContinuation> maybeContinuation = action.run(partition, record, tracker, watermarkEstimator);
    assertEquals(Optional.of(ProcessContinuation.stop()), maybeContinuation);
    verify(watermarkEstimator, never()).setWatermark(any());
    verify(dao, never()).insert(any());
}

Example 8 with PartitionMetadata

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

the class PartitionMetadataMapperTest method testMapPartitionMetadataFromResultSet.

@Test
public void testMapPartitionMetadataFromResultSet() {
    final Struct row = Struct.newBuilder().set(COLUMN_PARTITION_TOKEN).to("token").set(COLUMN_PARENT_TOKENS).toStringArray(Collections.singletonList("parentToken")).set(COLUMN_START_TIMESTAMP).to(Timestamp.ofTimeMicroseconds(10L)).set(COLUMN_END_TIMESTAMP).to(Timestamp.ofTimeMicroseconds(20L)).set(COLUMN_HEARTBEAT_MILLIS).to(5_000L).set(COLUMN_STATE).to(State.RUNNING.name()).set(COLUMN_WATERMARK).to(Timestamp.ofTimeMicroseconds(30L)).set(COLUMN_CREATED_AT).to(Timestamp.ofTimeMicroseconds(40L)).set(COLUMN_SCHEDULED_AT).to(Timestamp.ofTimeMicroseconds(50L)).set(COLUMN_RUNNING_AT).to(Timestamp.ofTimeMicroseconds(60L)).set(COLUMN_FINISHED_AT).to(Timestamp.ofTimeMicroseconds(70L)).build();
    final PartitionMetadata partition = mapper.from(row);
    assertEquals(new PartitionMetadata("token", Sets.newHashSet("parentToken"), Timestamp.ofTimeMicroseconds(10L), Timestamp.ofTimeMicroseconds(20L), 5_000L, State.RUNNING, Timestamp.ofTimeMicroseconds(30), Timestamp.ofTimeMicroseconds(40), Timestamp.ofTimeMicroseconds(50), Timestamp.ofTimeMicroseconds(60), Timestamp.ofTimeMicroseconds(70)), partition);
}

Example 9 with PartitionMetadata

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

the class ReadChangeStreamPartitionRangeTrackerTest method testTryClaim.

@Test
public void testTryClaim() {
    final PartitionMetadata partition = mock(PartitionMetadata.class);
    final OffsetRange range = new OffsetRange(100, 200);
    final ReadChangeStreamPartitionRangeTracker tracker = new ReadChangeStreamPartitionRangeTracker(partition, range);
    assertEquals(range, tracker.currentRestriction());
    assertTrue(tracker.tryClaim(100L));
    assertTrue(tracker.tryClaim(100L));
    assertTrue(tracker.tryClaim(150L));
    assertTrue(tracker.tryClaim(199L));
    assertFalse(tracker.tryClaim(200L));
}

Example 10 with PartitionMetadata

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