Search in sources :

Example 6 with PullRequest

use of com.google.pubsub.v1.PullRequest 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();
    }
}
Also used : StreamObserver(io.grpc.stub.StreamObserver) SubscriberImplBase(com.google.pubsub.v1.SubscriberGrpc.SubscriberImplBase) Server(io.grpc.Server) PullRequest(com.google.pubsub.v1.PullRequest) ArrayList(java.util.ArrayList) ByteString(com.google.protobuf.ByteString) ReceivedMessage(com.google.pubsub.v1.ReceivedMessage) Timestamp(com.google.protobuf.Timestamp) PubsubMessage(com.google.pubsub.v1.PubsubMessage) PullResponse(com.google.pubsub.v1.PullResponse) IncomingMessage(org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage) Test(org.junit.Test)

Example 7 with PullRequest

use of com.google.pubsub.v1.PullRequest 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();
    }
}
Also used : StreamObserver(io.grpc.stub.StreamObserver) SubscriberImplBase(com.google.pubsub.v1.SubscriberGrpc.SubscriberImplBase) Server(io.grpc.Server) PullRequest(com.google.pubsub.v1.PullRequest) ArrayList(java.util.ArrayList) ByteString(com.google.protobuf.ByteString) ReceivedMessage(com.google.pubsub.v1.ReceivedMessage) Timestamp(com.google.protobuf.Timestamp) PubsubMessage(com.google.pubsub.v1.PubsubMessage) PullResponse(com.google.pubsub.v1.PullResponse) IncomingMessage(org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage) Test(org.junit.Test)

Example 8 with PullRequest

use of com.google.pubsub.v1.PullRequest in project beam by apache.

the class PubsubGrpcClient method pull.

@Override
public List<IncomingMessage> pull(long requestTimeMsSinceEpoch, SubscriptionPath subscription, int batchSize, boolean returnImmediately) throws IOException {
    PullRequest request = PullRequest.newBuilder().setSubscription(subscription.getPath()).setReturnImmediately(returnImmediately).setMaxMessages(batchSize).build();
    PullResponse response = subscriberStub().pull(request);
    if (response.getReceivedMessagesCount() == 0) {
        return ImmutableList.of();
    }
    List<IncomingMessage> incomingMessages = new ArrayList<>(response.getReceivedMessagesCount());
    for (ReceivedMessage message : response.getReceivedMessagesList()) {
        PubsubMessage pubsubMessage = message.getMessage();
        @Nullable Map<String, String> attributes = pubsubMessage.getAttributes();
        // Timestamp.
        long timestampMsSinceEpoch;
        if (Strings.isNullOrEmpty(timestampAttribute)) {
            Timestamp timestampProto = pubsubMessage.getPublishTime();
            checkArgument(timestampProto != null, "Pubsub message is missing timestamp proto");
            timestampMsSinceEpoch = timestampProto.getSeconds() * 1000 + timestampProto.getNanos() / 1000L / 1000L;
        } 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 && attributes != null) {
            recordId = attributes.get(idAttribute);
        }
        if (Strings.isNullOrEmpty(recordId)) {
            // Fall back to the Pubsub provided message id.
            recordId = pubsubMessage.getMessageId();
        }
        incomingMessages.add(IncomingMessage.of(pubsubMessage, timestampMsSinceEpoch, requestTimeMsSinceEpoch, ackId, recordId));
    }
    return incomingMessages;
}
Also used : PullResponse(com.google.pubsub.v1.PullResponse) PullRequest(com.google.pubsub.v1.PullRequest) ArrayList(java.util.ArrayList) ReceivedMessage(com.google.pubsub.v1.ReceivedMessage) Timestamp(com.google.protobuf.Timestamp) PubsubMessage(com.google.pubsub.v1.PubsubMessage) Nullable(org.checkerframework.checker.nullness.qual.Nullable)

Example 9 with PullRequest

use of com.google.pubsub.v1.PullRequest in project flink by apache.

the class PubsubHelper method pullMessages.

// 
// Mostly copied from the example on https://cloud.google.com/pubsub/docs/pull
// Licensed under the Apache 2.0 License to "Google LLC" from
// https://github.com/googleapis/google-cloud-java/blob/master/google-cloud-examples/src/main/java/com/google/cloud/examples/pubsub/snippets/SubscriberSnippets.java.
// 
public List<ReceivedMessage> pullMessages(String projectId, String subscriptionId, int maxNumberOfMessages) throws Exception {
    SubscriberStubSettings subscriberStubSettings = SubscriberStubSettings.newBuilder().setTransportChannelProvider(channelProvider).setCredentialsProvider(EmulatorCredentialsProvider.create()).build();
    try (SubscriberStub subscriber = GrpcSubscriberStub.create(subscriberStubSettings)) {
        String subscriptionName = ProjectSubscriptionName.format(projectId, subscriptionId);
        PullRequest pullRequest = PullRequest.newBuilder().setMaxMessages(maxNumberOfMessages).setSubscription(subscriptionName).build();
        List<ReceivedMessage> receivedMessages = subscriber.pullCallable().call(pullRequest).getReceivedMessagesList();
        acknowledgeIds(subscriber, subscriptionName, receivedMessages);
        return receivedMessages;
    }
}
Also used : PullRequest(com.google.pubsub.v1.PullRequest) SubscriberStub(com.google.cloud.pubsub.v1.stub.SubscriberStub) GrpcSubscriberStub(com.google.cloud.pubsub.v1.stub.GrpcSubscriberStub) ReceivedMessage(com.google.pubsub.v1.ReceivedMessage) SubscriberStubSettings(com.google.cloud.pubsub.v1.stub.SubscriberStubSettings)

Aggregations

PullRequest (com.google.pubsub.v1.PullRequest)7 PullResponse (com.google.pubsub.v1.PullResponse)5 ReceivedMessage (com.google.pubsub.v1.ReceivedMessage)4 ByteString (com.google.protobuf.ByteString)3 Timestamp (com.google.protobuf.Timestamp)3 PubsubMessage (com.google.pubsub.v1.PubsubMessage)3 ArrayList (java.util.ArrayList)3 Test (org.junit.Test)3 SubscriberStub (com.google.cloud.pubsub.v1.stub.SubscriberStub)2 SubscriberGrpc (com.google.pubsub.v1.SubscriberGrpc)2 SubscriberImplBase (com.google.pubsub.v1.SubscriberGrpc.SubscriberImplBase)2 ManagedChannel (io.grpc.ManagedChannel)2 Server (io.grpc.Server)2 StreamObserver (io.grpc.stub.StreamObserver)2 IncomingMessage (org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage)2 GrpcSubscriberStub (com.google.cloud.pubsub.v1.stub.GrpcSubscriberStub)1 SubscriberStubSettings (com.google.cloud.pubsub.v1.stub.SubscriberStubSettings)1 GeneratedMessageV3 (com.google.protobuf.GeneratedMessageV3)1 AcknowledgeRequest (com.google.pubsub.v1.AcknowledgeRequest)1 StreamingPullRequest (com.google.pubsub.v1.StreamingPullRequest)1