Search in sources :

Example 1 with PullRequest

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

the class PubsubJsonClientTest method pullOneMessageEmptyAttributes.

@Test
public void pullOneMessageEmptyAttributes() throws IOException {
    client = new PubsubJsonClient(null, null, mockPubsub);
    String expectedSubscription = SUBSCRIPTION.getPath();
    PullRequest expectedRequest = new PullRequest().setReturnImmediately(true).setMaxMessages(10);
    PubsubMessage expectedPubsubMessage = new PubsubMessage().setMessageId(MESSAGE_ID).encodeData(DATA.getBytes(StandardCharsets.UTF_8)).setPublishTime(String.valueOf(PUB_TIME));
    ReceivedMessage expectedReceivedMessage = new ReceivedMessage().setMessage(expectedPubsubMessage).setAckId(ACK_ID);
    PullResponse expectedResponse = new PullResponse().setReceivedMessages(ImmutableList.of(expectedReceivedMessage));
    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, actualMessage.message().getData().toStringUtf8());
    assertEquals(REQ_TIME, actualMessage.requestTimeMsSinceEpoch());
    assertEquals(PUB_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) 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)

Example 2 with PullRequest

use of com.google.api.services.pubsub.model.PullRequest 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(StandardCharsets.UTF_8)).setPublishTime(String.valueOf(PUB_TIME)).setAttributes(ImmutableMap.of(TIMESTAMP_ATTRIBUTE, String.valueOf(MESSAGE_TIME), ID_ATTRIBUTE, RECORD_ID)).setOrderingKey(ORDERING_KEY);
    ReceivedMessage expectedReceivedMessage = new ReceivedMessage().setMessage(expectedPubsubMessage).setAckId(ACK_ID);
    PullResponse expectedResponse = new PullResponse().setReceivedMessages(ImmutableList.of(expectedReceivedMessage));
    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, actualMessage.message().getData().toStringUtf8());
    assertEquals(RECORD_ID, actualMessage.recordId());
    assertEquals(REQ_TIME, actualMessage.requestTimeMsSinceEpoch());
    assertEquals(MESSAGE_TIME, actualMessage.timestampMsSinceEpoch());
    assertEquals(ORDERING_KEY, actualMessage.message().getOrderingKey());
}
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)

Example 3 with PullRequest

use of com.google.api.services.pubsub.model.PullRequest 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 4 with PullRequest

use of com.google.api.services.pubsub.model.PullRequest 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

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 ReceivedMessage (com.google.api.services.pubsub.model.ReceivedMessage)4 ByteString (com.google.protobuf.ByteString)4 IncomingMessage (org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage)3 Test (org.junit.Test)3 ArrayList (java.util.ArrayList)1 Nullable (org.checkerframework.checker.nullness.qual.Nullable)1