Examples with PubsubReader org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubReader used on opensource projects

Example 6 with PubsubReader

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubReader 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();
}

Example 7 with PubsubReader

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubReader 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();
}

Example 8 with PubsubReader

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubUnboundedSource.PubsubReader 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));
}