use of com.google.pubsub.v1.PullRequest in project beam by apache.
the class PubsubGrpcClientTest method pullOneMessage.
@Test
public void pullOneMessage() throws IOException {
initializeClient(null, null);
String expectedSubscription = SUBSCRIPTION.getPath();
final PullRequest expectedRequest = PullRequest.newBuilder().setSubscription(expectedSubscription).setReturnImmediately(true).setMaxMessages(10).build();
Timestamp timestamp = Timestamp.newBuilder().setSeconds(PUB_TIME_MS / 1000).setNanos((int) (PUB_TIME_MS % 1000) * 1000 * 1000).build();
PubsubMessage expectedPubsubMessage = PubsubMessage.newBuilder().setMessageId(MESSAGE_ID).setData(ByteString.copyFrom(DATA.getBytes(StandardCharsets.UTF_8))).setPublishTime(timestamp).build();
ReceivedMessage expectedReceivedMessage = ReceivedMessage.newBuilder().setMessage(expectedPubsubMessage).setAckId(ACK_ID).build();
final PullResponse response = PullResponse.newBuilder().addAllReceivedMessages(ImmutableList.of(expectedReceivedMessage)).build();
final List<PullRequest> requestsReceived = new ArrayList<>();
SubscriberImplBase subscriberImplBase = new SubscriberImplBase() {
@Override
public void pull(PullRequest request, StreamObserver<PullResponse> responseObserver) {
requestsReceived.add(request);
responseObserver.onNext(response);
responseObserver.onCompleted();
}
};
Server server = InProcessServerBuilder.forName(channelName).addService(subscriberImplBase).build().start();
try {
List<IncomingMessage> actualMessages = client.pull(REQ_TIME_MS, SUBSCRIPTION, 10, true);
assertEquals(1, actualMessages.size());
IncomingMessage actualMessage = actualMessages.get(0);
assertEquals(ACK_ID, actualMessage.ackId());
assertEquals(DATA, actualMessage.message().getData().toStringUtf8());
assertEquals(MESSAGE_ID, actualMessage.recordId());
assertEquals(REQ_TIME_MS, actualMessage.requestTimeMsSinceEpoch());
assertEquals(PUB_TIME_MS, actualMessage.timestampMsSinceEpoch());
assertEquals(expectedRequest, Iterables.getOnlyElement(requestsReceived));
} finally {
server.shutdownNow();
}
}
use of com.google.pubsub.v1.PullRequest in project beam by apache.
the class PubsubGrpcClientTest method pullOneMessageUsingAttributes.
@Test
public void pullOneMessageUsingAttributes() throws IOException {
initializeClient(TIMESTAMP_ATTRIBUTE, ID_ATTRIBUTE);
String expectedSubscription = SUBSCRIPTION.getPath();
final PullRequest expectedRequest = PullRequest.newBuilder().setSubscription(expectedSubscription).setReturnImmediately(true).setMaxMessages(10).build();
Timestamp timestamp = Timestamp.newBuilder().setSeconds(PUB_TIME_MS / 1000).setNanos((int) (PUB_TIME_MS % 1000) * 1000 * 1000).build();
PubsubMessage expectedPubsubMessage = PubsubMessage.newBuilder().setMessageId(MESSAGE_ID).setData(ByteString.copyFrom(DATA.getBytes(StandardCharsets.UTF_8))).setPublishTime(timestamp).putAllAttributes(ATTRIBUTES).putAllAttributes(ImmutableMap.of(TIMESTAMP_ATTRIBUTE, String.valueOf(MESSAGE_TIME_MS), ID_ATTRIBUTE, RECORD_ID)).build();
ReceivedMessage expectedReceivedMessage = ReceivedMessage.newBuilder().setMessage(expectedPubsubMessage).setAckId(ACK_ID).build();
final PullResponse response = PullResponse.newBuilder().addAllReceivedMessages(ImmutableList.of(expectedReceivedMessage)).build();
final List<PullRequest> requestsReceived = new ArrayList<>();
SubscriberImplBase subscriberImplBase = new SubscriberImplBase() {
@Override
public void pull(PullRequest request, StreamObserver<PullResponse> responseObserver) {
requestsReceived.add(request);
responseObserver.onNext(response);
responseObserver.onCompleted();
}
};
Server server = InProcessServerBuilder.forName(channelName).addService(subscriberImplBase).build().start();
try {
List<IncomingMessage> actualMessages = client.pull(REQ_TIME_MS, SUBSCRIPTION, 10, true);
assertEquals(1, actualMessages.size());
IncomingMessage actualMessage = actualMessages.get(0);
assertEquals(ACK_ID, actualMessage.ackId());
assertEquals(DATA, actualMessage.message().getData().toStringUtf8());
assertEquals(RECORD_ID, actualMessage.recordId());
assertEquals(REQ_TIME_MS, actualMessage.requestTimeMsSinceEpoch());
assertEquals(MESSAGE_TIME_MS, actualMessage.timestampMsSinceEpoch());
assertEquals(expectedRequest, Iterables.getOnlyElement(requestsReceived));
} finally {
server.shutdownNow();
}
}
use of com.google.pubsub.v1.PullRequest in project beam by apache.
the class PubsubGrpcClient method pull.
@Override
public List<IncomingMessage> pull(long requestTimeMsSinceEpoch, SubscriptionPath subscription, int batchSize, boolean returnImmediately) throws IOException {
PullRequest request = PullRequest.newBuilder().setSubscription(subscription.getPath()).setReturnImmediately(returnImmediately).setMaxMessages(batchSize).build();
PullResponse response = subscriberStub().pull(request);
if (response.getReceivedMessagesCount() == 0) {
return ImmutableList.of();
}
List<IncomingMessage> incomingMessages = new ArrayList<>(response.getReceivedMessagesCount());
for (ReceivedMessage message : response.getReceivedMessagesList()) {
PubsubMessage pubsubMessage = message.getMessage();
@Nullable Map<String, String> attributes = pubsubMessage.getAttributes();
// Timestamp.
long timestampMsSinceEpoch;
if (Strings.isNullOrEmpty(timestampAttribute)) {
Timestamp timestampProto = pubsubMessage.getPublishTime();
checkArgument(timestampProto != null, "Pubsub message is missing timestamp proto");
timestampMsSinceEpoch = timestampProto.getSeconds() * 1000 + timestampProto.getNanos() / 1000L / 1000L;
} else {
timestampMsSinceEpoch = extractTimestampAttribute(timestampAttribute, attributes);
}
// Ack id.
String ackId = message.getAckId();
checkState(!Strings.isNullOrEmpty(ackId));
// Record id, if any.
@Nullable String recordId = null;
if (idAttribute != null && attributes != null) {
recordId = attributes.get(idAttribute);
}
if (Strings.isNullOrEmpty(recordId)) {
// Fall back to the Pubsub provided message id.
recordId = pubsubMessage.getMessageId();
}
incomingMessages.add(IncomingMessage.of(pubsubMessage, timestampMsSinceEpoch, requestTimeMsSinceEpoch, ackId, recordId));
}
return incomingMessages;
}
use of com.google.pubsub.v1.PullRequest in project flink by apache.
the class PubsubHelper method pullMessages.
//
// Mostly copied from the example on https://cloud.google.com/pubsub/docs/pull
// Licensed under the Apache 2.0 License to "Google LLC" from
// https://github.com/googleapis/google-cloud-java/blob/master/google-cloud-examples/src/main/java/com/google/cloud/examples/pubsub/snippets/SubscriberSnippets.java.
//
public List<ReceivedMessage> pullMessages(String projectId, String subscriptionId, int maxNumberOfMessages) throws Exception {
SubscriberStubSettings subscriberStubSettings = SubscriberStubSettings.newBuilder().setTransportChannelProvider(channelProvider).setCredentialsProvider(EmulatorCredentialsProvider.create()).build();
try (SubscriberStub subscriber = GrpcSubscriberStub.create(subscriberStubSettings)) {
String subscriptionName = ProjectSubscriptionName.format(projectId, subscriptionId);
PullRequest pullRequest = PullRequest.newBuilder().setMaxMessages(maxNumberOfMessages).setSubscription(subscriptionName).build();
List<ReceivedMessage> receivedMessages = subscriber.pullCallable().call(pullRequest).getReceivedMessagesList();
acknowledgeIds(subscriber, subscriptionName, receivedMessages);
return receivedMessages;
}
}
Aggregations