Search in sources :

Example 1 with PartitionMetadataDao

use of org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataDao in project beam by apache.

the class InitializeDoFn method processElement.

@ProcessElement
public void processElement(OutputReceiver<PartitionMetadata> receiver) {
    PartitionMetadataDao partitionMetadataDao = daoFactory.getPartitionMetadataDao();
    if (!partitionMetadataDao.tableExists()) {
        daoFactory.getPartitionMetadataAdminDao().createPartitionMetadataTable();
        createFakeParentPartition();
    }
    final PartitionMetadata initialPartition = Optional.ofNullable(partitionMetadataDao.getPartition(InitialPartition.PARTITION_TOKEN)).map(mapperFactory.partitionMetadataMapper()::from).orElseThrow(() -> new IllegalStateException("Initial partition not found in metadata table."));
    receiver.output(initialPartition);
}
Also used : PartitionMetadataDao(org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataDao) PartitionMetadata(org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionMetadata)

Example 2 with PartitionMetadataDao

use of org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataDao in project beam by apache.

the class QueryChangeStreamActionTest method setUp.

@Before
public void setUp() throws Exception {
    changeStreamDao = mock(ChangeStreamDao.class);
    partitionMetadataDao = mock(PartitionMetadataDao.class);
    changeStreamRecordMapper = mock(ChangeStreamRecordMapper.class);
    partitionMetadataMapper = mock(PartitionMetadataMapper.class);
    dataChangeRecordAction = mock(DataChangeRecordAction.class);
    heartbeatRecordAction = mock(HeartbeatRecordAction.class);
    childPartitionsRecordAction = mock(ChildPartitionsRecordAction.class);
    action = new QueryChangeStreamAction(changeStreamDao, partitionMetadataDao, changeStreamRecordMapper, partitionMetadataMapper, dataChangeRecordAction, heartbeatRecordAction, childPartitionsRecordAction);
    final Struct row = mock(Struct.class);
    partition = PartitionMetadata.newBuilder().setPartitionToken(PARTITION_TOKEN).setParentTokens(Sets.newHashSet("parentToken")).setStartTimestamp(PARTITION_START_TIMESTAMP).setEndTimestamp(PARTITION_END_TIMESTAMP).setHeartbeatMillis(PARTITION_HEARTBEAT_MILLIS).setState(SCHEDULED).setWatermark(WATERMARK_TIMESTAMP).setScheduledAt(Timestamp.now()).build();
    restriction = mock(OffsetRange.class);
    restrictionTracker = mock(RestrictionTracker.class);
    outputReceiver = mock(OutputReceiver.class);
    watermarkEstimator = mock(ManualWatermarkEstimator.class);
    bundleFinalizer = new BundleFinalizerStub();
    when(restrictionTracker.currentRestriction()).thenReturn(restriction);
    when(restriction.getFrom()).thenReturn(10L);
    when(partitionMetadataDao.getPartition(PARTITION_TOKEN)).thenReturn(row);
    when(partitionMetadataMapper.from(row)).thenReturn(partition);
}
Also used : RestrictionTracker(org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker) ChangeStreamDao(org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.ChangeStreamDao) ChangeStreamRecordMapper(org.apache.beam.sdk.io.gcp.spanner.changestreams.mapper.ChangeStreamRecordMapper) OutputReceiver(org.apache.beam.sdk.transforms.DoFn.OutputReceiver) Struct(com.google.cloud.spanner.Struct) OffsetRange(org.apache.beam.sdk.io.range.OffsetRange) PartitionMetadataDao(org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataDao) PartitionMetadataMapper(org.apache.beam.sdk.io.gcp.spanner.changestreams.mapper.PartitionMetadataMapper) ManualWatermarkEstimator(org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator) Before(org.junit.Before)

Example 3 with PartitionMetadataDao

use of org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataDao in project beam by apache.

the class ReadChangeStreamPartitionDoFnTest method setUp.

@Before
public void setUp() {
    final DaoFactory daoFactory = mock(DaoFactory.class);
    final MapperFactory mapperFactory = mock(MapperFactory.class);
    final ChangeStreamMetrics metrics = mock(ChangeStreamMetrics.class);
    final ActionFactory actionFactory = mock(ActionFactory.class);
    final PartitionMetadataDao partitionMetadataDao = mock(PartitionMetadataDao.class);
    final ChangeStreamDao changeStreamDao = mock(ChangeStreamDao.class);
    final ChangeStreamRecordMapper changeStreamRecordMapper = mock(ChangeStreamRecordMapper.class);
    final PartitionMetadataMapper partitionMetadataMapper = mock(PartitionMetadataMapper.class);
    dataChangeRecordAction = mock(DataChangeRecordAction.class);
    heartbeatRecordAction = mock(HeartbeatRecordAction.class);
    childPartitionsRecordAction = mock(ChildPartitionsRecordAction.class);
    queryChangeStreamAction = mock(QueryChangeStreamAction.class);
    doFn = new ReadChangeStreamPartitionDoFn(daoFactory, mapperFactory, actionFactory, metrics);
    partition = PartitionMetadata.newBuilder().setPartitionToken(PARTITION_TOKEN).setParentTokens(Sets.newHashSet("parentToken")).setStartTimestamp(PARTITION_START_TIMESTAMP).setEndTimestamp(PARTITION_END_TIMESTAMP).setHeartbeatMillis(PARTITION_HEARTBEAT_MILLIS).setState(SCHEDULED).setWatermark(PARTITION_START_TIMESTAMP).setScheduledAt(Timestamp.now()).build();
    restriction = mock(OffsetRange.class);
    restrictionTracker = mock(RestrictionTracker.class);
    outputReceiver = mock(OutputReceiver.class);
    watermarkEstimator = mock(ManualWatermarkEstimator.class);
    bundleFinalizer = mock(BundleFinalizer.class);
    when(restrictionTracker.currentRestriction()).thenReturn(restriction);
    when(daoFactory.getPartitionMetadataDao()).thenReturn(partitionMetadataDao);
    when(daoFactory.getChangeStreamDao()).thenReturn(changeStreamDao);
    when(mapperFactory.changeStreamRecordMapper()).thenReturn(changeStreamRecordMapper);
    when(mapperFactory.partitionMetadataMapper()).thenReturn(partitionMetadataMapper);
    when(actionFactory.dataChangeRecordAction()).thenReturn(dataChangeRecordAction);
    when(actionFactory.heartbeatRecordAction(metrics)).thenReturn(heartbeatRecordAction);
    when(actionFactory.childPartitionsRecordAction(partitionMetadataDao, metrics)).thenReturn(childPartitionsRecordAction);
    when(actionFactory.queryChangeStreamAction(changeStreamDao, partitionMetadataDao, changeStreamRecordMapper, partitionMetadataMapper, dataChangeRecordAction, heartbeatRecordAction, childPartitionsRecordAction)).thenReturn(queryChangeStreamAction);
    doFn.setup();
}
Also used : RestrictionTracker(org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker) ChangeStreamDao(org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.ChangeStreamDao) ChangeStreamRecordMapper(org.apache.beam.sdk.io.gcp.spanner.changestreams.mapper.ChangeStreamRecordMapper) ChildPartitionsRecordAction(org.apache.beam.sdk.io.gcp.spanner.changestreams.action.ChildPartitionsRecordAction) DataChangeRecordAction(org.apache.beam.sdk.io.gcp.spanner.changestreams.action.DataChangeRecordAction) QueryChangeStreamAction(org.apache.beam.sdk.io.gcp.spanner.changestreams.action.QueryChangeStreamAction) HeartbeatRecordAction(org.apache.beam.sdk.io.gcp.spanner.changestreams.action.HeartbeatRecordAction) OutputReceiver(org.apache.beam.sdk.transforms.DoFn.OutputReceiver) ChangeStreamMetrics(org.apache.beam.sdk.io.gcp.spanner.changestreams.ChangeStreamMetrics) DaoFactory(org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.DaoFactory) BundleFinalizer(org.apache.beam.sdk.transforms.DoFn.BundleFinalizer) OffsetRange(org.apache.beam.sdk.io.range.OffsetRange) PartitionMetadataDao(org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataDao) PartitionMetadataMapper(org.apache.beam.sdk.io.gcp.spanner.changestreams.mapper.PartitionMetadataMapper) MapperFactory(org.apache.beam.sdk.io.gcp.spanner.changestreams.mapper.MapperFactory) ActionFactory(org.apache.beam.sdk.io.gcp.spanner.changestreams.action.ActionFactory) ManualWatermarkEstimator(org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator) Before(org.junit.Before)

Example 4 with PartitionMetadataDao

use of org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataDao in project beam by apache.

the class DetectNewPartitionsDoFn method setup.

/**
 * Obtains the instance of {@link DetectNewPartitionsAction}.
 */
@Setup
public void setup() {
    final PartitionMetadataDao partitionMetadataDao = daoFactory.getPartitionMetadataDao();
    final PartitionMetadataMapper partitionMetadataMapper = mapperFactory.partitionMetadataMapper();
    this.detectNewPartitionsAction = actionFactory.detectNewPartitionsAction(partitionMetadataDao, partitionMetadataMapper, metrics, resumeDuration);
}
Also used : PartitionMetadataDao(org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataDao) PartitionMetadataMapper(org.apache.beam.sdk.io.gcp.spanner.changestreams.mapper.PartitionMetadataMapper)

Example 5 with PartitionMetadataDao

use of org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataDao in project beam by apache.

the class ReadChangeStreamPartitionDoFn method setup.

/**
 * Constructs instances for the {@link PartitionMetadataDao}, {@link ChangeStreamDao}, {@link
 * ChangeStreamRecordMapper}, {@link PartitionMetadataMapper}, {@link DataChangeRecordAction},
 * {@link HeartbeatRecordAction}, {@link ChildPartitionsRecordAction} and {@link
 * QueryChangeStreamAction}.
 */
@Setup
public void setup() {
    final PartitionMetadataDao partitionMetadataDao = daoFactory.getPartitionMetadataDao();
    final ChangeStreamDao changeStreamDao = daoFactory.getChangeStreamDao();
    final ChangeStreamRecordMapper changeStreamRecordMapper = mapperFactory.changeStreamRecordMapper();
    final PartitionMetadataMapper partitionMetadataMapper = mapperFactory.partitionMetadataMapper();
    final DataChangeRecordAction dataChangeRecordAction = actionFactory.dataChangeRecordAction();
    final HeartbeatRecordAction heartbeatRecordAction = actionFactory.heartbeatRecordAction(metrics);
    final ChildPartitionsRecordAction childPartitionsRecordAction = actionFactory.childPartitionsRecordAction(partitionMetadataDao, metrics);
    this.queryChangeStreamAction = actionFactory.queryChangeStreamAction(changeStreamDao, partitionMetadataDao, changeStreamRecordMapper, partitionMetadataMapper, dataChangeRecordAction, heartbeatRecordAction, childPartitionsRecordAction);
}
Also used : ChangeStreamDao(org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.ChangeStreamDao) PartitionMetadataDao(org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataDao) ChangeStreamRecordMapper(org.apache.beam.sdk.io.gcp.spanner.changestreams.mapper.ChangeStreamRecordMapper) PartitionMetadataMapper(org.apache.beam.sdk.io.gcp.spanner.changestreams.mapper.PartitionMetadataMapper) ChildPartitionsRecordAction(org.apache.beam.sdk.io.gcp.spanner.changestreams.action.ChildPartitionsRecordAction) DataChangeRecordAction(org.apache.beam.sdk.io.gcp.spanner.changestreams.action.DataChangeRecordAction) HeartbeatRecordAction(org.apache.beam.sdk.io.gcp.spanner.changestreams.action.HeartbeatRecordAction)

Aggregations

PartitionMetadataDao (org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.PartitionMetadataDao)5 PartitionMetadataMapper (org.apache.beam.sdk.io.gcp.spanner.changestreams.mapper.PartitionMetadataMapper)4 ChangeStreamDao (org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.ChangeStreamDao)3 ChangeStreamRecordMapper (org.apache.beam.sdk.io.gcp.spanner.changestreams.mapper.ChangeStreamRecordMapper)3 ChildPartitionsRecordAction (org.apache.beam.sdk.io.gcp.spanner.changestreams.action.ChildPartitionsRecordAction)2 DataChangeRecordAction (org.apache.beam.sdk.io.gcp.spanner.changestreams.action.DataChangeRecordAction)2 HeartbeatRecordAction (org.apache.beam.sdk.io.gcp.spanner.changestreams.action.HeartbeatRecordAction)2 OffsetRange (org.apache.beam.sdk.io.range.OffsetRange)2 OutputReceiver (org.apache.beam.sdk.transforms.DoFn.OutputReceiver)2 ManualWatermarkEstimator (org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator)2 RestrictionTracker (org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker)2 Before (org.junit.Before)2 Struct (com.google.cloud.spanner.Struct)1 ChangeStreamMetrics (org.apache.beam.sdk.io.gcp.spanner.changestreams.ChangeStreamMetrics)1 ActionFactory (org.apache.beam.sdk.io.gcp.spanner.changestreams.action.ActionFactory)1 QueryChangeStreamAction (org.apache.beam.sdk.io.gcp.spanner.changestreams.action.QueryChangeStreamAction)1 DaoFactory (org.apache.beam.sdk.io.gcp.spanner.changestreams.dao.DaoFactory)1 MapperFactory (org.apache.beam.sdk.io.gcp.spanner.changestreams.mapper.MapperFactory)1 PartitionMetadata (org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionMetadata)1 BundleFinalizer (org.apache.beam.sdk.transforms.DoFn.BundleFinalizer)1