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());
}
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;
}
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));
}
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);
}
}
}
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);
}
}
}
}
Aggregations