Search in sources :

Example 6 with PubsubCheckpoint

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubCheckpoint in project beam by apache.

the class PubsubUnboundedSourceTest method readOneMessage.

@Test
public void readOneMessage() throws IOException {
    setupOneMessage();
    PubsubReader reader = primSource.createReader(p.getOptions(), null);
    // Read one message.
    assertTrue(reader.start());
    assertEquals(DATA, data(reader.getCurrent()));
    assertFalse(reader.advance());
    // ACK the message.
    PubsubCheckpoint checkpoint = reader.getCheckpointMark();
    checkpoint.finalizeCheckpoint();
    reader.close();
}
Also used : PubsubCheckpoint(org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubCheckpoint) PubsubReader(org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubReader) Test(org.junit.Test)

Example 7 with PubsubCheckpoint

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubCheckpoint in project beam by apache.

the class PubsubUnboundedSourceTest method timeoutAckExtensions.

@Test
public void timeoutAckExtensions() throws IOException {
    setupOneMessage();
    PubsubReader reader = primSource.createReader(p.getOptions(), null);
    PubsubTestClient pubsubClient = (PubsubTestClient) reader.getPubsubClient();
    // Pull the first message but don't take a checkpoint for it.
    assertTrue(reader.start());
    assertEquals(DATA, data(reader.getCurrent()));
    // Extend the ack.
    now.addAndGet(55 * 1000);
    pubsubClient.advance();
    assertFalse(reader.advance());
    // Let the ack expire.
    for (int i = 0; i < 3; i++) {
        now.addAndGet(25 * 1000);
        pubsubClient.advance();
        assertFalse(reader.advance());
    }
    // Wait for resend.
    now.addAndGet(25 * 1000);
    pubsubClient.advance();
    // Reread the same message.
    assertTrue(reader.advance());
    assertEquals(DATA, data(reader.getCurrent()));
    // Now ACK the message.
    PubsubCheckpoint checkpoint = reader.getCheckpointMark();
    checkpoint.finalizeCheckpoint();
    reader.close();
}
Also used : PubsubCheckpoint(org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubCheckpoint) PubsubReader(org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubReader) PubsubCheckpoint(org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubCheckpoint) Test(org.junit.Test)

Example 8 with PubsubCheckpoint

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubCheckpoint in project beam by apache.

the class PubsubUnboundedSourceTest method noSubscriptionNoSplitGeneratesSubscription.

@Test
public void noSubscriptionNoSplitGeneratesSubscription() throws Exception {
    TopicPath topicPath = PubsubClient.topicPathFromName("my_project", "my_topic");
    factory = PubsubTestClient.createFactoryForCreateSubscription();
    PubsubUnboundedSource source = new PubsubUnboundedSource(factory, StaticValueProvider.of(PubsubClient.projectPathFromId("my_project")), StaticValueProvider.of(topicPath), null, /* subscription */
    null, /* timestampLabel */
    null, /* idLabel */
    false);
    assertThat(source.getSubscription(), nullValue());
    assertThat(source.getSubscription(), nullValue());
    PipelineOptions options = PipelineOptionsFactory.create();
    PubsubSource actualSource = new PubsubSource(source);
    PubsubReader reader = actualSource.createReader(options, null);
    SubscriptionPath createdSubscription = reader.subscription;
    assertThat(createdSubscription, not(nullValue()));
    PubsubCheckpoint checkpoint = reader.getCheckpointMark();
    assertThat(checkpoint.subscriptionPath, equalTo(createdSubscription.getPath()));
    checkpoint.finalizeCheckpoint();
    PubsubCheckpoint deserCheckpoint = CoderUtils.clone(actualSource.getCheckpointMarkCoder(), checkpoint);
    assertThat(checkpoint.subscriptionPath, not(nullValue()));
    assertThat(checkpoint.subscriptionPath, equalTo(deserCheckpoint.subscriptionPath));
    PubsubReader readerFromOriginal = actualSource.createReader(options, checkpoint);
    PubsubReader readerFromDeser = actualSource.createReader(options, deserCheckpoint);
    assertThat(readerFromOriginal.subscription, equalTo(createdSubscription));
    assertThat(readerFromDeser.subscription, equalTo(createdSubscription));
}
Also used : TopicPath(org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath) PipelineOptions(org.apache.beam.sdk.options.PipelineOptions) SubscriptionPath(org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath) PubsubSource(org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubSource) PubsubCheckpoint(org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubCheckpoint) PubsubReader(org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubReader) Test(org.junit.Test)

Aggregations

PubsubCheckpoint (org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubCheckpoint)8 PubsubReader (org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubReader)8 Test (org.junit.Test)8 ArrayList (java.util.ArrayList)2 IncomingMessage (org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage)2 HashMap (java.util.HashMap)1 SubscriptionPath (org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath)1 TopicPath (org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath)1 PubsubSource (org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubSource)1 PipelineOptions (org.apache.beam.sdk.options.PipelineOptions)1 Instant (org.joda.time.Instant)1