Search in sources :

Example 1 with PullResponse

use of com.google.api.services.pubsub.model.PullResponse in project beam by apache.

the class PubsubJsonClientTest method pullOneMessage.

@Test
public void pullOneMessage() throws IOException {
    String expectedSubscription = SUBSCRIPTION.getPath();
    PullRequest expectedRequest = new PullRequest().setReturnImmediately(true).setMaxMessages(10);
    PubsubMessage expectedPubsubMessage = new PubsubMessage().setMessageId(MESSAGE_ID).encodeData(DATA.getBytes()).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);
    assertEquals(ACK_ID, actualMessage.ackId);
    assertEquals(DATA, new String(actualMessage.elementBytes));
    assertEquals(RECORD_ID, actualMessage.recordId);
    assertEquals(REQ_TIME, actualMessage.requestTimeMsSinceEpoch);
    assertEquals(MESSAGE_TIME, actualMessage.timestampMsSinceEpoch);
}
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) ReceivedMessage(com.google.api.services.pubsub.model.ReceivedMessage) PubsubMessage(com.google.api.services.pubsub.model.PubsubMessage) Test(org.junit.Test)

Example 2 with PullResponse

use of com.google.api.services.pubsub.model.PullResponse in project beam by apache.

the class PubsubJsonClient method pull.

@Override
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().size() == 0) {
        return ImmutableList.of();
    }
    List<IncomingMessage> incomingMessages = new ArrayList<>(response.getReceivedMessages().size());
    for (ReceivedMessage message : response.getReceivedMessages()) {
        PubsubMessage pubsubMessage = message.getMessage();
        @Nullable Map<String, String> attributes = pubsubMessage.getAttributes();
        // Payload.
        byte[] elementBytes = pubsubMessage.decodeData();
        // Timestamp.
        long timestampMsSinceEpoch = extractTimestamp(timestampAttribute, message.getMessage().getPublishTime(), 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(new IncomingMessage(elementBytes, attributes, timestampMsSinceEpoch, requestTimeMsSinceEpoch, ackId, recordId));
    }
    return incomingMessages;
}
Also used : PullResponse(com.google.api.services.pubsub.model.PullResponse) PullRequest(com.google.api.services.pubsub.model.PullRequest) ArrayList(java.util.ArrayList) ReceivedMessage(com.google.api.services.pubsub.model.ReceivedMessage) PubsubMessage(com.google.api.services.pubsub.model.PubsubMessage) Nullable(javax.annotation.Nullable)

Aggregations

PubsubMessage (com.google.api.services.pubsub.model.PubsubMessage)2 PullRequest (com.google.api.services.pubsub.model.PullRequest)2 PullResponse (com.google.api.services.pubsub.model.PullResponse)2 ReceivedMessage (com.google.api.services.pubsub.model.ReceivedMessage)2 ArrayList (java.util.ArrayList)1 Nullable (javax.annotation.Nullable)1 IncomingMessage (org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage)1 Test (org.junit.Test)1