use of com.google.api.services.pubsub.model.ReceivedMessage in project beam by apache.
the class PubsubJsonClient method pull.
@Override
@SuppressWarnings("ProtoFieldNullComparison")
public List<IncomingMessage> pull(long requestTimeMsSinceEpoch, SubscriptionPath subscription, int batchSize, boolean returnImmediately) throws IOException {
PullRequest request = new PullRequest().setReturnImmediately(returnImmediately).setMaxMessages(batchSize);
PullResponse response = pubsub.projects().subscriptions().pull(subscription.getPath(), request).execute();
if (response.getReceivedMessages() == null || response.getReceivedMessages().isEmpty()) {
return ImmutableList.of();
}
List<IncomingMessage> incomingMessages = new ArrayList<>(response.getReceivedMessages().size());
for (ReceivedMessage message : response.getReceivedMessages()) {
PubsubMessage pubsubMessage = message.getMessage();
Map<String, String> attributes;
if (pubsubMessage.getAttributes() != null) {
attributes = pubsubMessage.getAttributes();
} else {
attributes = new HashMap<>();
}
// Payload.
byte[] elementBytes = pubsubMessage.getData() == null ? null : pubsubMessage.decodeData();
if (elementBytes == null) {
elementBytes = new byte[0];
}
// Timestamp.
long timestampMsSinceEpoch;
if (Strings.isNullOrEmpty(timestampAttribute)) {
timestampMsSinceEpoch = parseTimestampAsMsSinceEpoch(message.getMessage().getPublishTime());
} 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) {
recordId = attributes.get(idAttribute);
}
if (Strings.isNullOrEmpty(recordId)) {
// Fall back to the Pubsub provided message id.
recordId = pubsubMessage.getMessageId();
}
com.google.pubsub.v1.PubsubMessage.Builder protoMessage = com.google.pubsub.v1.PubsubMessage.newBuilder();
protoMessage.setData(ByteString.copyFrom(elementBytes));
protoMessage.putAllAttributes(attributes);
// PubsubMessage uses `null` to represent no ordering key where we want a default of "".
if (pubsubMessage.getOrderingKey() != null) {
protoMessage.setOrderingKey(pubsubMessage.getOrderingKey());
} else {
protoMessage.setOrderingKey("");
}
incomingMessages.add(IncomingMessage.of(protoMessage.build(), timestampMsSinceEpoch, requestTimeMsSinceEpoch, ackId, recordId));
}
return incomingMessages;
}
use of com.google.api.services.pubsub.model.ReceivedMessage in project beam by apache.
the class PubsubJsonClientTest method pullOneMessageWithNoData.
@Test
public void pullOneMessageWithNoData() throws IOException {
String expectedSubscription = SUBSCRIPTION.getPath();
PullRequest expectedRequest = new PullRequest().setReturnImmediately(true).setMaxMessages(10);
PubsubMessage expectedPubsubMessage = new PubsubMessage().setMessageId(MESSAGE_ID).setPublishTime(String.valueOf(PUB_TIME)).setAttributes(ImmutableMap.of(TIMESTAMP_ATTRIBUTE, String.valueOf(MESSAGE_TIME), ID_ATTRIBUTE, RECORD_ID));
ReceivedMessage expectedReceivedMessage = new ReceivedMessage().setMessage(expectedPubsubMessage).setAckId(ACK_ID);
PullResponse expectedResponse = new PullResponse().setReceivedMessages(ImmutableList.of(expectedReceivedMessage));
Mockito.when((Object) mockPubsub.projects().subscriptions().pull(expectedSubscription, expectedRequest).execute()).thenReturn(expectedResponse);
List<IncomingMessage> acutalMessages = client.pull(REQ_TIME, SUBSCRIPTION, 10, true);
assertEquals(1, acutalMessages.size());
IncomingMessage actualMessage = acutalMessages.get(0);
assertArrayEquals(new byte[0], actualMessage.message().getData().toByteArray());
}
Aggregations