Search in sources :

Example 6 with DataChangeRecord

use of org.apache.beam.sdk.io.gcp.spanner.changestreams.model.DataChangeRecord 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());
}
Also used : ChangeStreamResultSet(org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.ChangeStreamResultSet) ChangeStreamResultSetMetadata(org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.ChangeStreamResultSetMetadata) DataChangeRecord(org.apache.beam.sdk.io.gcp.spanner.changestreams.model.DataChangeRecord) Struct(com.google.cloud.spanner.Struct) ProcessContinuation(org.apache.beam.sdk.transforms.DoFn.ProcessContinuation) Test(org.junit.Test)

Example 7 with DataChangeRecord

use of org.apache.beam.sdk.io.gcp.spanner.changestreams.model.DataChangeRecord 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);
}
Also used : Instant(org.joda.time.Instant) DataChangeRecord(org.apache.beam.sdk.io.gcp.spanner.changestreams.model.DataChangeRecord) Timestamp(com.google.cloud.Timestamp) ProcessContinuation(org.apache.beam.sdk.transforms.DoFn.ProcessContinuation) Test(org.junit.Test)

Aggregations

DataChangeRecord (org.apache.beam.sdk.io.gcp.spanner.changestreams.model.DataChangeRecord)7 Test (org.junit.Test)6 Struct (com.google.cloud.spanner.Struct)4 ProcessContinuation (org.apache.beam.sdk.transforms.DoFn.ProcessContinuation)4 Timestamp (com.google.cloud.Timestamp)3 ColumnType (org.apache.beam.sdk.io.gcp.spanner.changestreams.model.ColumnType)3 Mod (org.apache.beam.sdk.io.gcp.spanner.changestreams.model.Mod)3 TypeCode (org.apache.beam.sdk.io.gcp.spanner.changestreams.model.TypeCode)3 ChangeStreamResultSet (org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.ChangeStreamResultSet)2 SpannerException (com.google.cloud.spanner.SpannerException)1 Scope (io.opencensus.common.Scope)1 ChangeStreamResultSetMetadata (org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.ChangeStreamResultSetMetadata)1 ChangeStreamRecord (org.apache.beam.sdk.io.gcp.spanner.changestreams.model.ChangeStreamRecord)1 ChildPartitionsRecord (org.apache.beam.sdk.io.gcp.spanner.changestreams.model.ChildPartitionsRecord)1 HeartbeatRecord (org.apache.beam.sdk.io.gcp.spanner.changestreams.model.HeartbeatRecord)1 PartitionMetadata (org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionMetadata)1 VisibleForTesting (org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting)1 Instant (org.joda.time.Instant)1