Examples with IncomingMessage org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage used on opensource projects

Example 1 with IncomingMessage

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage 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();
    }
}

Example 2 with IncomingMessage

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage 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();
    }
}

Example 3 with IncomingMessage

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage 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());
}

Example 4 with IncomingMessage

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage 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());
}

Example 5 with IncomingMessage

use of org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage in project beam by apache.

the class PubsubUnboundedSourceTest method readManyMessages.

@Test
public void readManyMessages() throws Exception {
    Map<String, Integer> dataToMessageNum = new HashMap<>();
    final int m = 97;
    final int n = 10000;
    List<IncomingMessage> incoming = new ArrayList<>();
    for (int i = 0; i < n; i++) {
        // Make the messages timestamps slightly out of order.
        int messageNum = ((i / m) * m) + (m - 1) - (i % m);
        String data = String.format("data_%d", messageNum);
        dataToMessageNum.put(data, messageNum);
        String recid = String.format("recordid_%d", messageNum);
        String ackId = String.format("ackid_%d", messageNum);
        incoming.add(IncomingMessage.of(com.google.pubsub.v1.PubsubMessage.newBuilder().setData(ByteString.copyFromUtf8(data)).build(), messageNumToTimestamp(messageNum), 0, ackId, recid));
    }
    setupOneMessage(incoming);
    PubsubReader reader = primSource.createReader(p.getOptions(), null);
    PubsubTestClient pubsubClient = (PubsubTestClient) reader.getPubsubClient();
    for (int i = 0; i < n; i++) {
        if (i == 0) {
            assertTrue(reader.start());
        } else {
            assertTrue(reader.advance());
        }
        // We'll checkpoint and ack within the 2min limit.
        now.addAndGet(30);
        pubsubClient.advance();
        String data = data(reader.getCurrent(), !(primSource.outer.getNeedsAttributes() || primSource.outer.getNeedsMessageId()));
        Integer messageNum = dataToMessageNum.remove(data);
        // No duplicate messages.
        assertNotNull(messageNum);
        // Preserve timestamp.
        assertEquals(new Instant(messageNumToTimestamp(messageNum)), reader.getCurrentTimestamp());
        // Preserve record id.
        String recid = String.format("recordid_%d", messageNum);
        assertArrayEquals(recid.getBytes(StandardCharsets.UTF_8), reader.getCurrentRecordId());
        if (i % 1000 == 999) {
            // Estimated watermark can never get ahead of actual outstanding messages.
            long watermark = reader.getWatermark().getMillis();
            long minOutstandingTimestamp = Long.MAX_VALUE;
            for (Integer outstandingMessageNum : dataToMessageNum.values()) {
                minOutstandingTimestamp = Math.min(minOutstandingTimestamp, messageNumToTimestamp(outstandingMessageNum));
            }
            assertThat(watermark, lessThanOrEqualTo(minOutstandingTimestamp));
            // Ack messages, but only every other finalization.
            PubsubCheckpoint checkpoint = reader.getCheckpointMark();
            if (i % 2000 == 1999) {
                checkpoint.finalizeCheckpoint();
            }
        }
    }
    // We are done.
    assertFalse(reader.advance());
    // We saw each message exactly once.
    assertTrue(dataToMessageNum.isEmpty());
    reader.close();
}