com.google.pubsub.v1.PubsubMessage Java Examples

private List<String> receiveLines(int expectedMessageCount) throws Exception {
  List<String> received = new CopyOnWriteArrayList<>();
  ProjectSubscriptionName subscriptionName = ProjectSubscriptionName.of(projectId,
      subscriptionId);

  MessageReceiver receiver = ((PubsubMessage message, AckReplyConsumer consumer) -> {
    try {
      String encoded = Json.asString(new org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage(
          message.getData().toByteArray(), message.getAttributesMap()));
      received.add(encoded);
    } catch (IOException e) {
      throw new UncheckedIOException(e);
    }
    consumer.ack();
  });
  Subscriber subscriber = Subscriber.newBuilder(subscriptionName, receiver).build();
  subscriber.startAsync();
  while (received.size() < expectedMessageCount) {
    Thread.sleep(100);
  }
  subscriber.stopAsync();

  return received;
}
 

@Override
public void send(PubsubMessage message) throws Exception {
  String wadoUri = message.getData().toStringUtf8();
  String qidoUri = qidoFromWadoUri(wadoUri);

  // Invoke QIDO-RS to get DICOM tags needed to invoke C-Store.
  JSONArray qidoResponse = dicomWebClient.qidoRs(qidoUri);
  if (qidoResponse.length() != 1) {
    throw new IllegalArgumentException(
        "Invalid QidoRS JSON array length for response: " + qidoResponse.toString());
  }
  String sopClassUid = AttributesUtil.getTagValue(qidoResponse.getJSONObject(0),
      SOP_CLASS_UID_TAG);
  String sopInstanceUid = AttributesUtil.getTagValue(qidoResponse.getJSONObject(0),
      SOP_INSTANCE_UID_TAG);

  // Invoke WADO-RS to get bulk DICOM.
  InputStream responseStream = dicomWebClient.wadoRs(wadoUri);

  CountingInputStream countingStream = new CountingInputStream(responseStream);
  DicomClient.connectAndCstore(sopClassUid, sopInstanceUid, countingStream,
      applicationEntity, dimsePeerAET, dimsePeerIP, dimsePeerPort);
  MonitoringService.addEvent(Event.BYTES, countingStream.getCount());
}
 

@Override
public String apply(PubsubMessage message) {
  Map<String, String> attributes = new HashMap<>(message.getAttributesMap());

  // We coerce all docType and namespace names to be snake_case and to remove invalid
  // characters; these transformations MUST match with the transformations applied by the
  // jsonschema-transpiler and mozilla-schema-generator when creating table schemas in BigQuery.
  final String namespace = attributes.get(Attribute.DOCUMENT_NAMESPACE);
  final String docType = attributes.get(Attribute.DOCUMENT_TYPE);
  if (namespace != null) {
    attributes.put(Attribute.DOCUMENT_NAMESPACE, getAndCacheNormalizedName(namespace));
  }
  if (docType != null) {
    attributes.put(Attribute.DOCUMENT_TYPE, getAndCacheNormalizedName(docType));
  }

  attributes = Maps.transformValues(DerivedAttributesMap.of(attributes),
      v -> v.replaceAll("-", "_"));

  return super.apply(PubsubMessage.newBuilder().putAllAttributes(attributes).build());
}
 

@Override
public ListenableFuture<List<PubsubMessage>> pullAndAckAsync(String subscription, Integer maxMessages,
		Boolean returnImmediately) {
	PullRequest pullRequest = this.subscriberFactory.createPullRequest(
			subscription, maxMessages, returnImmediately);

	final SettableListenableFuture<List<PubsubMessage>> settableFuture = new SettableListenableFuture<>();

	this.pullAsync(pullRequest).addCallback(
			ackableMessages -> {
				if (!ackableMessages.isEmpty()) {
					ack(ackableMessages);
				}
				List<PubsubMessage> messages = ackableMessages.stream()
						.map(AcknowledgeablePubsubMessage::getPubsubMessage)
						.collect(Collectors.toList());

				settableFuture.set(messages);
			},
			settableFuture::setException);

	return settableFuture;
}
 

@Override
public <T> T fromPubSubMessage(PubsubMessage message, Class<T> payloadType) {
	T result;
	byte[] payload = message.getData().toByteArray();

	if (payloadType == ByteString.class) {
		result = (T) message.getData();
	}
	else if (payloadType == String.class) {
		result = (T) new String(payload, this.charset);
	}
	else if (payloadType == ByteBuffer.class) {
		result = (T) ByteBuffer.wrap(payload);
	}
	else if (payloadType == byte[].class) {
		result = (T) payload;
	}
	else {
		throw new PubSubMessageConversionException("Unable to convert Pub/Sub message to payload of type " +
				payloadType.getName() + ".");
	}

	return result;
}
 

void processMessage(SourceContext<OUT> sourceContext, List<ReceivedMessage> messages) throws Exception {
	synchronized (sourceContext.getCheckpointLock()) {
		for (ReceivedMessage message : messages) {
			acknowledgeOnCheckpoint.addAcknowledgeId(message.getAckId());

			PubsubMessage pubsubMessage = message.getMessage();

			OUT deserializedMessage = deserializationSchema.deserialize(pubsubMessage);
			if (deserializationSchema.isEndOfStream(deserializedMessage)) {
				cancel();
				return;
			}

			sourceContext.collect(deserializedMessage);
		}

	}
}
 

private void parseMessage(PubsubMessage message) {
  LOG.trace("parseMessage called");
  try (DataParser parser = parserFactory.getParser(message.getMessageId(), message.getData().toByteArray())) {
    Record r;
    while ((r = tryParse(parser, message)) != null) {
      setHeaders(message, r);

      batchMaker.addRecord(r);
      ++currentRecordCount;
    }
  } catch (DataParserException | IOException e) {
    LOG.error(Errors.PUBSUB_05.getMessage(), e.toString(), e);
    // Create a raw record of the gRPC message data, set attributes as headers, and use the messageId for the recordId
    Record errorRecord = context.createRecord(message.getMessageId(), message.getData().toByteArray(), MIME_GRPC);
    setHeaders(message, errorRecord);
    context.reportError(new OnRecordErrorException(errorRecord, Errors.PUBSUB_05, e.toString()));
  }
}
 

/**
 * Tests that the correct message is sent to the publisher when the record has a null value.
 */
@Test
public void testPutWithNullValues() {
  props.put(CloudPubSubSinkConnector.MAX_BUFFER_SIZE_CONFIG, CPS_MIN_BATCH_SIZE1);
  task.start(props);
  List<SinkRecord> records = new ArrayList<>();
  records.add(
      new SinkRecord(
          KAFKA_TOPIC,
          0,
          STRING_SCHEMA,
          KAFKA_MESSAGE_KEY,
          STRING_SCHEMA,
          null,
          -1));
  task.put(records);
  ArgumentCaptor<PubsubMessage> captor = ArgumentCaptor.forClass(PubsubMessage.class);
  verify(publisher, times(1)).publish(captor.capture());
  List<PubsubMessage> requestArgs = captor.getAllValues();
  List<PubsubMessage> expectedMessages = new ArrayList<>();
  Map<String, String> attributes = new HashMap<>();
  attributes.put(ConnectorUtils.CPS_MESSAGE_KEY_ATTRIBUTE, KAFKA_MESSAGE_KEY);
  expectedMessages.add(
      PubsubMessage.newBuilder().putAllAttributes(attributes).build());
  assertEquals(requestArgs, expectedMessages);
}
 

@Test
public void testMessagesAreAbandonedOnNonRetriableFailure() throws IOException {
    // Simulate a failure on send that indicates a retry isn't allowed.
    final Publisher publisher = mockPublisher.orElseThrow(IllegalStateException::new);
    when(publisher.publish(any(PubsubMessage.class)))
        .thenReturn(failedFuture(new ApiException("simulated permanent failure",
                                                  new IOException(),
                                                  GrpcStatusCode.of(Status.Code.NOT_FOUND),
                                                  false)));
    // Here we send the message.
    processSingleMessage();

    // Now we check the invocations…
    verify(publisher).publish(any(PubsubMessage.class));
    verifyNoMoreInteractions(publisher);
}
 

@Override
public void doPost(HttpServletRequest req, HttpServletResponse resp)
    throws IOException, ServletException {
  Publisher publisher = this.publisher;
  try {
    String topicId = System.getenv("PUBSUB_TOPIC");
    // create a publisher on the topic
    if (publisher == null) {
      publisher = Publisher.newBuilder(
          ProjectTopicName.of(ServiceOptions.getDefaultProjectId(), topicId))
          .build();
    }
    // construct a pubsub message from the payload
    final String payload = req.getParameter("payload");
    PubsubMessage pubsubMessage =
        PubsubMessage.newBuilder().setData(ByteString.copyFromUtf8(payload)).build();

    publisher.publish(pubsubMessage);
    // redirect to home page
    resp.sendRedirect("/");
  } catch (Exception e) {
    resp.sendError(HttpStatus.SC_INTERNAL_SERVER_ERROR, e.getMessage());
  }
}
 

@Override
public ListenableFuture<Void> publish(
    int clientId, int sequenceNumber, long publishTimestampMillis) {
  SettableFuture<Void> done = SettableFuture.create();
  ApiFutures.addCallback(
      publisher.publish(
          PubsubMessage.newBuilder()
              .setData(payload)
              .putAttributes("sendTime", Long.toString(publishTimestampMillis))
              .putAttributes("clientId", Integer.toString(clientId))
              .putAttributes("sequenceNumber", Integer.toString(sequenceNumber))
              .build()),
      new ApiFutureCallback<String>() {
        @Override
        public void onSuccess(String messageId) {
          done.set(null);
        }

        @Override
        public void onFailure(Throwable t) {
          done.setException(t);
        }
      },
      MoreExecutors.directExecutor());
  return done;
}
 

@Test
public void testPull() throws Exception {
	Publisher publisher = pubsubHelper.createPublisher(PROJECT_NAME, TOPIC_NAME);
	publisher
		.publish(PubsubMessage
			.newBuilder()
			.setData(ByteString.copyFromUtf8("Hello World PULL"))
			.build())
		.get();

	List<ReceivedMessage> receivedMessages = pubsubHelper.pullMessages(PROJECT_NAME, SUBSCRIPTION_NAME, 10);
	assertEquals(1, receivedMessages.size());
	assertEquals("Hello World PULL", receivedMessages.get(0).getMessage().getData().toStringUtf8());

	publisher.shutdown();
}
 

@Override
public CompletableFuture<String> apply(PubsubMessage message) {
  final PubsubMessage compressed = compress.apply(message);
  final ApiFuture<String> future = getPublisher(message).publish(compressed);
  final CompletableFuture<String> result = new CompletableFuture<>();
  ApiFutures.addCallback(future, new ApiFutureCallback<String>() {

    @Override
    public void onFailure(Throwable throwable) {
      result.completeExceptionally(throwable);
    }

    @Override
    public void onSuccess(String messageId) {
      result.complete(messageId);
    }
  }, executor);
  return result;
}
 

private static void handleMessage(
    DlpJob job,
    SettableApiFuture<Boolean> done,
    PubsubMessage pubsubMessage,
    AckReplyConsumer ackReplyConsumer) {
  String messageAttribute = pubsubMessage.getAttributesMap().get("DlpJobName");
  if (job.getName().equals(messageAttribute)) {
    done.set(true);
    ackReplyConsumer.ack();
  } else {
    ackReplyConsumer.nack();
  }
}
 

private static void handleMessage(
    DlpJob job,
    SettableApiFuture<Boolean> done,
    PubsubMessage pubsubMessage,
    AckReplyConsumer ackReplyConsumer) {
  String messageAttribute = pubsubMessage.getAttributesMap().get("DlpJobName");
  if (job.getName().equals(messageAttribute)) {
    done.set(true);
    ackReplyConsumer.ack();
  } else {
    ackReplyConsumer.nack();
  }
}
 

@Override
public boolean test(PubsubMessage message) {
  String docType = template.apply(message);
  try {
    return cache.get(docType, () -> pattern.matcher(docType.replaceAll("_", "-")).matches());
  } catch (ExecutionException e) {
    throw new UncheckedExecutionException(e.getCause());
  }
}
 

@Test
public void canLimitBatchByteSize() {
  PubsubMessage input = PubsubMessage.newBuilder().putAttributes("meta", "data").build();
  for (int i = 0; i < Math.ceil((MAX_BYTES + 1) / input.getSerializedSize()); i++) {
    output.apply(input);
  }
  assertThat((int) output.batches.get(BATCH_KEY).byteSize, lessThanOrEqualTo(MAX_BYTES));
}
 

@Test
public void canWriteToStaticDestination() {
  new Pubsub.Write(pubsub.getTopic(), ForkJoinPool.commonPool(), b -> b, m -> m)
      .apply(PubsubMessage.newBuilder()
          .setData(ByteString.copyFrom("test".getBytes(StandardCharsets.UTF_8))).build())
      .join();
  assertEquals(ImmutableList.of("test"), pubsub.pull(1, false).stream()
      .map(m -> m.getData().toStringUtf8()).collect(Collectors.toList()));
}
 

@Test
public void testString() throws JSONException {
	String str = "test 123";

	PubsubMessage pubsubMessage = this.converter.toPubSubMessage(str, null);

	JSONAssert.assertEquals(
			"\"test 123\"",
			pubsubMessage.getData().toStringUtf8(),
			true);

	Object o = this.converter.fromPubSubMessage(pubsubMessage, String.class);

	assertThat(o).as("verify that deserialized object is equal to the original one").isEqualTo(str);
}
 

/** Generate a sequence of PubsubMessage objects. */
static List<PubsubMessage> generatePubsubMessages(int howMany) {
  List<PubsubMessage> messages = new ArrayList<>();
  for (int i = 0; i < howMany; i++) {
    messages.add(
        PubsubMessage.newBuilder().setData(ByteString.copyFrom("message-" + i, UTF_8)).build());
  }
  return messages;
}
 

@Override
public void invoke(IN message, SinkFunction.Context context) throws Exception {
	PubsubMessage pubsubMessage = PubsubMessage
		.newBuilder()
		.setData(ByteString.copyFrom(serializationSchema.serialize(message)))
		.build();

	ApiFuture<String> future = publisher.publish(pubsubMessage);
	numPendingFutures.incrementAndGet();
	ApiFutures.addCallback(future, failureHandler, directExecutor());
}
 

@Override
public void service(HttpRequest request, HttpResponse response) throws IOException {
  Optional<String> maybeTopicName = request.getFirstQueryParameter("topic");
  Optional<String> maybeMessage = request.getFirstQueryParameter("message");

  BufferedWriter responseWriter = response.getWriter();

  if (maybeTopicName.isEmpty() || maybeMessage.isEmpty()) {
    response.setStatusCode(HttpURLConnection.HTTP_BAD_REQUEST);

    responseWriter.write("Missing 'topic' and/or 'subscription' parameter(s).");
    return;
  }

  String topicName = maybeTopicName.get();
  logger.info("Publishing message to topic: " + topicName);

  // Create the PubsubMessage object
  // (This is different than the PubSubMessage POJO used in Pub/Sub-triggered functions)
  ByteString byteStr = ByteString.copyFrom(maybeMessage.get(), StandardCharsets.UTF_8);
  PubsubMessage pubsubApiMessage = PubsubMessage.newBuilder().setData(byteStr).build();

  Publisher publisher = Publisher.newBuilder(
      ProjectTopicName.of(PROJECT_ID, topicName)).build();

  // Attempt to publish the message
  String responseMessage;
  try {
    publisher.publish(pubsubApiMessage).get();
    responseMessage = "Message published.";
  } catch (InterruptedException | ExecutionException e) {
    logger.log(Level.SEVERE, "Error publishing Pub/Sub message: " + e.getMessage(), e);
    responseMessage = "Error publishing Pub/Sub message; see logs for more info.";
  }

  responseWriter.write(responseMessage);
}
 

@Override
public void receiveMessage(PubsubMessage message, AckReplyConsumer consumer) {
  try {
    messages.put(new MessageReplyConsumerBundle(message, consumer));
  } catch (InterruptedException e) {
    LOG.warn(
        "Thread interrupted while trying to enqueue message with id '{}'. Sending nack. Message will be re-received",
        message.getMessageId()
    );
    consumer.nack();
    Thread.currentThread().interrupt();
  }
}
 

/**
 * Tests that a call to flush() processes the Futures that were generated by calls to put.
 */
@Test
public void testFlushWithNoPublishInPut() throws Exception {
  task.start(props);
  Map<TopicPartition, OffsetAndMetadata> partitionOffsets = new HashMap<>();
  partitionOffsets.put(new TopicPartition(KAFKA_TOPIC, 0), null);
  List<SinkRecord> records = getSampleRecords();
  ApiFuture<String> goodFuture = getSuccessfulPublishFuture();
  when(publisher.publish(any(PubsubMessage.class))).thenReturn(goodFuture);
  task.put(records);
  task.flush(partitionOffsets);
  verify(publisher, times(2)).publish(any(PubsubMessage.class));
  verify(goodFuture, times(2)).addListener(any(Runnable.class), any(Executor.class));
}
 

@Test
public void modifyAckDeadline() {
  int partitions = 3;
  int recordsPerPartition = 2;
  generateTestRecordsForConsumers(partitions, recordsPerPartition, null);

  // Pull 3 times to get all records
  List<String> ackIds = new ArrayList<>();
  for (int i = 0; i < partitions; i++) {
    for (PubsubMessage message : subscriptionManager.pull(10, false)) {
      ackIds.add(message.getMessageId());
    }
  }
  assertEquals(ackIds, subscriptionManager.modifyAckDeadline(ackIds, 10));

  // Should be after the initial expiration
  when(mockClock.instant()).thenReturn(FIXED_INSTANT.plusSeconds(15));

  assertEquals(ackIds, subscriptionManager.acknowledge(ackIds));

  subscriptionManager.runOneIteration();
  for (int i = 0; i < partitions; i++) {
    TopicPartition topicPartition = new TopicPartition(KAFKA_TOPIC_NAME, i);
    assertEquals(
        2,
        kafkaClientFactory
            .getConsumersForSubscription(SUBSCRIPTION.getName())
            .get(i)
            .position(topicPartition));
  }
}
 

@Test
public void testCreatePublishPullNextAndDelete() {
	pubSubTest((PubSubTemplate pubSubTemplate, String subscriptionName, String topicName) -> {
		//tag::publish[]
		Map<String, String> headers = Collections.singletonMap("key1", "val1");
		pubSubTemplate.publish(topicName, "message", headers).get();
		//end::publish[]
		PubsubMessage pubsubMessage = pubSubTemplate.pullNext(subscriptionName);

		assertThat(pubsubMessage.getData()).isEqualTo(ByteString.copyFromUtf8("message"));
		assertThat(pubsubMessage.getAttributesCount()).isEqualTo(1);
		assertThat(pubsubMessage.getAttributesOrThrow("key1")).isEqualTo("val1");
	});
}
 

@Test
public void canHandleDynamicTableIdWithEmptyValues() {
  output = new BigQuery.Write(bigQuery, MAX_BYTES, MAX_MESSAGES, NO_DELAY,
      PubsubMessageToTemplatedString.forBigQuery("${dataset}_.${table}_"),
      ForkJoinPool.commonPool(), PubsubMessageToObjectNode.Raw.of());
  output.apply(
      PubsubMessage.newBuilder().putAttributes("dataset", "").putAttributes("table", "").build())
      .join();
  assertNotNull(output.batches.get(BATCH_KEY));
}
 

@Test
public void pull_withHeader() {
  int partitions = 1;
  int recordsPerPartition = 3;
  List<Header> headers = new ArrayList<>();
  headers.add(new RecordHeader("key1", "value1".getBytes()));
  headers.add(new RecordHeader("key2", "value2".getBytes()));
  generateTestRecordsForConsumers(partitions, recordsPerPartition, headers);

  // Each response should pull from a different partition
  List<String> messageIds = new ArrayList<>();
  List<String> messages = new ArrayList<>();
  List<Map<String, String>> attributes = new ArrayList<>();
  List<PubsubMessage> response = subscriptionManager.pull(10, false);
  for (PubsubMessage message : response) {
    messageIds.add(message.getMessageId());
    messages.add(message.getData().toStringUtf8());
    attributes.add(message.getAttributesMap());
  }

  assertThat(messageIds, Matchers.contains("0-0", "0-1", "0-2"));
  assertThat(messages, Matchers.contains("message-0000", "message-0001", "message-0002"));
  ImmutableMap<String, String> expectedAttributes =
      new Builder<String, String>().put("key1", "value1").put("key2", "value2").build();
  assertThat(
      attributes,
      Matchers.equalTo(
          Arrays.asList(expectedAttributes, expectedAttributes, expectedAttributes)));

  assertThat(subscriptionManager.pull(10, false), Matchers.empty());
}
 

/** Constructor. */
public Write(String topicTemplate, Executor executor,
    Function<Publisher.Builder, Publisher.Builder> config,
    Function<PubsubMessage, PubsubMessage> compress) {
  this.executor = executor;
  this.topicTemplate = PubsubMessageToTemplatedString.of(topicTemplate);
  this.config = config;
  this.compress = compress;
}
 

public String sendMessage(Map<String, String> attributes, String body) {
  try {
    PubsubMessage message = PubsubMessage.newBuilder()
        .setData(ByteString.copyFrom(body, Charset.defaultCharset()))
        .putAllAttributes(attributes)
        .build();
    ApiFuture<String> publish = publisher.publish(message);
    return publish.get();
  } catch (Exception e) {
    throw new RuntimeException("While sending to topic " + registrar_topic, e);
  }
}