Search in sources :

Example 6 with ReceivedMessage

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;
}
Also used : PullRequest(com.google.api.services.pubsub.model.PullRequest) ArrayList(java.util.ArrayList) ByteString(com.google.protobuf.ByteString) ReceivedMessage(com.google.api.services.pubsub.model.ReceivedMessage) PubsubMessage(com.google.api.services.pubsub.model.PubsubMessage) PullResponse(com.google.api.services.pubsub.model.PullResponse) Nullable(org.checkerframework.checker.nullness.qual.Nullable)

Example 7 with ReceivedMessage

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());
}
Also used : PullResponse(com.google.api.services.pubsub.model.PullResponse) IncomingMessage(org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage) PullRequest(com.google.api.services.pubsub.model.PullRequest) ByteString(com.google.protobuf.ByteString) ReceivedMessage(com.google.api.services.pubsub.model.ReceivedMessage) PubsubMessage(com.google.api.services.pubsub.model.PubsubMessage) Test(org.junit.Test)

Aggregations

ReceivedMessage (com.google.api.services.pubsub.model.ReceivedMessage)7 PubsubMessage (com.google.api.services.pubsub.model.PubsubMessage)4 PullRequest (com.google.api.services.pubsub.model.PullRequest)4 PullResponse (com.google.api.services.pubsub.model.PullResponse)4 ByteString (com.google.protobuf.ByteString)4 ArrayList (java.util.ArrayList)4 Random (java.util.Random)3 IndexedRecord (org.apache.avro.generic.IndexedRecord)3 IncomingMessage (org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage)3 PCollection (org.apache.beam.sdk.values.PCollection)3 Test (org.junit.Test)3 ConvertToIndexedRecord (org.talend.components.adapter.beam.transform.ConvertToIndexedRecord)3 PTransform (org.apache.beam.sdk.transforms.PTransform)2 Nullable (org.checkerframework.checker.nullness.qual.Nullable)1 PubSubOutputProperties (org.talend.components.pubsub.output.PubSubOutputProperties)1