Search in sources :

Example 6 with SubscriptionPath

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath in project beam by apache.

the class TestPubsub method waitForNMessages.

/**
 * Repeatedly pull messages from {@link #subscriptionPath()}, returns after receiving {@code n}
 * messages or after waiting for {@code timeoutDuration}.
 */
public List<PubsubMessage> waitForNMessages(int n, Duration timeoutDuration) throws IOException, InterruptedException {
    Preconditions.checkNotNull(subscriptionPath);
    BlockingQueue<com.google.pubsub.v1.PubsubMessage> receivedMessages = new LinkedBlockingDeque<>(n);
    MessageReceiver receiver = (com.google.pubsub.v1.PubsubMessage message, AckReplyConsumer replyConsumer) -> {
        if (receivedMessages.offer(message)) {
            replyConsumer.ack();
        } else {
            replyConsumer.nack();
        }
    };
    Subscriber subscriber = Subscriber.newBuilder(subscriptionPath.getPath(), receiver).setCredentialsProvider(pipelineOptions::getGcpCredential).setChannelProvider(channelProvider).setEndpoint(pubsubEndpoint).build();
    subscriber.startAsync();
    DateTime startTime = new DateTime();
    int timeoutSeconds = timeoutDuration.toStandardSeconds().getSeconds();
    while (receivedMessages.size() < n && Seconds.secondsBetween(startTime, new DateTime()).getSeconds() < timeoutSeconds) {
        try {
            Thread.sleep(1000);
        } catch (InterruptedException ignored) {
        }
    }
    subscriber.stopAsync();
    subscriber.awaitTerminated();
    return receivedMessages.stream().map((message) -> new PubsubMessage(message.getData().toByteArray(), message.getAttributesMap(), message.getMessageId())).collect(Collectors.toList());
}
Also used : Statement(org.junit.runners.model.Statement) TopicPath(org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath) TestPipelineOptions(org.apache.beam.sdk.testing.TestPipelineOptions) ManagedChannel(io.grpc.ManagedChannel) TestRule(org.junit.rules.TestRule) Duration(org.joda.time.Duration) TimeoutException(java.util.concurrent.TimeoutException) Subscriber(com.google.cloud.pubsub.v1.Subscriber) Streams(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Streams) AckReplyConsumer(com.google.cloud.pubsub.v1.AckReplyConsumer) Publisher(com.google.cloud.pubsub.v1.Publisher) Seconds(org.joda.time.Seconds) FixedTransportChannelProvider(com.google.api.gax.rpc.FixedTransportChannelProvider) ThreadLocalRandom(java.util.concurrent.ThreadLocalRandom) TestPipeline(org.apache.beam.sdk.testing.TestPipeline) Iterables(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables) SubscriptionAdminClient(com.google.cloud.pubsub.v1.SubscriptionAdminClient) MatcherAssert.assertThat(org.hamcrest.MatcherAssert.assertThat) PipelineOptions(org.apache.beam.sdk.options.PipelineOptions) Nullable(org.checkerframework.checker.nullness.qual.Nullable) DateTimeFormat(org.joda.time.format.DateTimeFormat) MessageReceiver(com.google.cloud.pubsub.v1.MessageReceiver) ApiFutures(com.google.api.core.ApiFutures) DateTimeFormatter(org.joda.time.format.DateTimeFormatter) SubscriptionPath(org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath) DateTime(org.joda.time.DateTime) IOException(java.io.IOException) BlockingQueue(java.util.concurrent.BlockingQueue) Description(org.junit.runner.Description) TopicAdminSettings(com.google.cloud.pubsub.v1.TopicAdminSettings) Collectors(java.util.stream.Collectors) ApiFuture(com.google.api.core.ApiFuture) ByteString(com.google.protobuf.ByteString) ExecutionException(java.util.concurrent.ExecutionException) GrpcTransportChannel(com.google.api.gax.grpc.GrpcTransportChannel) ManagedChannelBuilder(io.grpc.ManagedChannelBuilder) List(java.util.List) Matchers.containsInAnyOrder(org.hamcrest.Matchers.containsInAnyOrder) TransportChannelProvider(com.google.api.gax.rpc.TransportChannelProvider) TopicAdminClient(com.google.cloud.pubsub.v1.TopicAdminClient) Preconditions(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions) Matcher(org.hamcrest.Matcher) Instant(org.joda.time.Instant) LinkedBlockingDeque(java.util.concurrent.LinkedBlockingDeque) SubscriptionAdminSettings(com.google.cloud.pubsub.v1.SubscriptionAdminSettings) PushConfig(com.google.pubsub.v1.PushConfig) LinkedBlockingDeque(java.util.concurrent.LinkedBlockingDeque) DateTime(org.joda.time.DateTime) MessageReceiver(com.google.cloud.pubsub.v1.MessageReceiver) Subscriber(com.google.cloud.pubsub.v1.Subscriber) AckReplyConsumer(com.google.cloud.pubsub.v1.AckReplyConsumer)

Example 7 with SubscriptionPath

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath in project beam by apache.

the class TestPubsub method initializePubsub.

private void initializePubsub(Description description) throws IOException {
    if (isLocalhost) {
        channel = ManagedChannelBuilder.forTarget(pubsubEndpoint).usePlaintext().build();
    } else {
        channel = ManagedChannelBuilder.forTarget(pubsubEndpoint).useTransportSecurity().build();
    }
    channelProvider = FixedTransportChannelProvider.create(GrpcTransportChannel.create(channel));
    topicAdmin = TopicAdminClient.create(TopicAdminSettings.newBuilder().setCredentialsProvider(pipelineOptions::getGcpCredential).setTransportChannelProvider(channelProvider).setEndpoint(pubsubEndpoint).build());
    subscriptionAdmin = SubscriptionAdminClient.create(SubscriptionAdminSettings.newBuilder().setCredentialsProvider(pipelineOptions::getGcpCredential).setTransportChannelProvider(channelProvider).setEndpoint(pubsubEndpoint).build());
    TopicPath eventsTopicPathTmp = PubsubClient.topicPathFromName(pipelineOptions.getProject(), createTopicName(description, EVENTS_TOPIC_NAME));
    topicAdmin.createTopic(eventsTopicPathTmp.getPath());
    // Set this after successful creation; it signals that the topic needs teardown
    eventsTopicPath = eventsTopicPathTmp;
    String subscriptionName = topicPath().getName() + "_beam_" + ThreadLocalRandom.current().nextLong();
    SubscriptionPath subscriptionPathTmp = new SubscriptionPath(String.format("projects/%s/subscriptions/%s", pipelineOptions.getProject(), subscriptionName));
    subscriptionAdmin.createSubscription(subscriptionPathTmp.getPath(), topicPath().getPath(), PushConfig.getDefaultInstance(), DEFAULT_ACK_DEADLINE_SECONDS);
    subscriptionPath = subscriptionPathTmp;
}
Also used : TopicPath(org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath) SubscriptionPath(org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath) ByteString(com.google.protobuf.ByteString)

Example 8 with SubscriptionPath

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath 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)

Example 9 with SubscriptionPath

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath in project beam by apache.

the class PubsubHelper method cleanup.

/**
 * Delete all the subscriptions and topics we created.
 */
public void cleanup() {
    for (SubscriptionPath subscription : createdSubscriptions) {
        try {
            NexmarkUtils.console("delete subscription %s", subscription);
            pubsubClient.deleteSubscription(subscription);
        } catch (IOException ex) {
            NexmarkUtils.console("could not delete subscription %s", subscription);
        }
    }
    for (TopicPath topic : createdTopics) {
        try {
            NexmarkUtils.console("delete topic %s", topic);
            pubsubClient.deleteTopic(topic);
        } catch (IOException ex) {
            NexmarkUtils.console("could not delete topic %s", topic);
        }
    }
}
Also used : TopicPath(org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath) SubscriptionPath(org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath) IOException(java.io.IOException)

Example 10 with SubscriptionPath

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath in project beam by apache.

the class PubsubHelper method createSubscription.

/**
 * Create subscription from short name. Ensure the subscription will be deleted on cleanup. Return
 * full subscription name.
 */
public SubscriptionPath createSubscription(String shortTopic, String shortSubscription) throws IOException {
    TopicPath topic = PubsubClient.topicPathFromName(project, shortTopic);
    SubscriptionPath subscription = PubsubClient.subscriptionPathFromName(project, shortSubscription);
    while (true) {
        try {
            NexmarkUtils.console("create subscription %s", subscription);
            pubsubClient.createSubscription(topic, subscription, 60);
            createdSubscriptions.add(subscription);
            return subscription;
        } catch (GoogleJsonResponseException ex) {
            NexmarkUtils.console("attempting to cleanup subscription %s", subscription);
            pubsubClient.deleteSubscription(subscription);
            try {
                if (!BackOffUtils.next(sleeper, backOff)) {
                    NexmarkUtils.console("too many retries for creating subscription %s", subscription);
                    throw ex;
                }
            } catch (InterruptedException in) {
                throw new IOException(in);
            }
        }
    }
}
Also used : GoogleJsonResponseException(com.google.api.client.googleapis.json.GoogleJsonResponseException) TopicPath(org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath) SubscriptionPath(org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath) IOException(java.io.IOException)

Aggregations

SubscriptionPath (org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath)13 TopicPath (org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath)10 IOException (java.io.IOException)5 Test (org.junit.Test)3 ApiException (com.google.api.gax.rpc.ApiException)2 ByteString (com.google.protobuf.ByteString)2 PipelineOptions (org.apache.beam.sdk.options.PipelineOptions)2 GoogleJsonResponseException (com.google.api.client.googleapis.json.GoogleJsonResponseException)1 ApiFuture (com.google.api.core.ApiFuture)1 ApiFutures (com.google.api.core.ApiFutures)1 GrpcTransportChannel (com.google.api.gax.grpc.GrpcTransportChannel)1 FixedTransportChannelProvider (com.google.api.gax.rpc.FixedTransportChannelProvider)1 TransportChannelProvider (com.google.api.gax.rpc.TransportChannelProvider)1 Subscriptions (com.google.api.services.pubsub.Pubsub.Projects.Subscriptions)1 ListSubscriptionsResponse (com.google.api.services.pubsub.model.ListSubscriptionsResponse)1 AckReplyConsumer (com.google.cloud.pubsub.v1.AckReplyConsumer)1 MessageReceiver (com.google.cloud.pubsub.v1.MessageReceiver)1 Publisher (com.google.cloud.pubsub.v1.Publisher)1 Subscriber (com.google.cloud.pubsub.v1.Subscriber)1 SubscriptionAdminClient (com.google.cloud.pubsub.v1.SubscriptionAdminClient)1