Examples with PubSubMessage com.google.cloud.pubsublite.proto.PubSubMessage used on opensource projects

Example 1 with PubSubMessage

use of com.google.cloud.pubsublite.proto.PubSubMessage in project beam by apache.

the class RowHandlerTest method fullMessageToRow.

@Test
public void fullMessageToRow() {
    RowHandler rowHandler = new RowHandler(FULL_READ_SCHEMA);
    Instant event = Instant.now();
    Instant publish = Instant.now();
    PubSubMessage userMessage = PubSubMessage.newBuilder().setKey(ByteString.copyFromUtf8("val1")).setData(ByteString.copyFromUtf8("val2")).setEventTime(Timestamps.fromMillis(event.getMillis())).putAttributes("key1", AttributeValues.newBuilder().addValues(ByteString.copyFromUtf8("attr1")).addValues(ByteString.copyFromUtf8("attr2")).build()).putAttributes("key2", AttributeValues.newBuilder().addValues(ByteString.copyFromUtf8("attr3")).build()).build();
    SequencedMessage sequencedMessage = SequencedMessage.newBuilder().setMessage(userMessage).setPublishTime(Timestamps.fromMillis(publish.getMillis())).build();
    Row expected = Row.withSchema(FULL_READ_SCHEMA).withFieldValue(RowHandler.MESSAGE_KEY_FIELD, "val1".getBytes(UTF_8)).withFieldValue(RowHandler.PAYLOAD_FIELD, "val2".getBytes(UTF_8)).withFieldValue(RowHandler.EVENT_TIMESTAMP_FIELD, event).withFieldValue(RowHandler.PUBLISH_TIMESTAMP_FIELD, publish).withFieldValue(RowHandler.ATTRIBUTES_FIELD, ImmutableList.of(Row.withSchema(RowHandler.ATTRIBUTES_ENTRY_SCHEMA).attachValues("key1", ImmutableList.of("attr1".getBytes(UTF_8), "attr2".getBytes(UTF_8))), Row.withSchema(RowHandler.ATTRIBUTES_ENTRY_SCHEMA).attachValues("key2", ImmutableList.of("attr3".getBytes(UTF_8))))).build();
    assertEquals(expected, rowHandler.messageToRow(sequencedMessage));
}

Example 2 with PubSubMessage

use of com.google.cloud.pubsublite.proto.PubSubMessage in project beam by apache.

the class RowHandlerTest method reorderRowToMessage.

@Test
public void reorderRowToMessage() {
    Schema schema = Schema.builder().addByteArrayField(RowHandler.MESSAGE_KEY_FIELD).addByteArrayField(RowHandler.PAYLOAD_FIELD).build();
    Schema rowSchema = Schema.builder().addByteArrayField(RowHandler.PAYLOAD_FIELD).addByteArrayField(RowHandler.MESSAGE_KEY_FIELD).build();
    RowHandler rowHandler = new RowHandler(schema);
    Row row = Row.withSchema(rowSchema).attachValues("abc".getBytes(UTF_8), "def".getBytes(UTF_8));
    PubSubMessage expected = PubSubMessage.newBuilder().setData(ByteString.copyFromUtf8("abc")).setKey(ByteString.copyFromUtf8("def")).build();
    assertEquals(expected, rowHandler.rowToMessage(row));
}

Example 3 with PubSubMessage

use of com.google.cloud.pubsublite.proto.PubSubMessage in project beam by apache.

the class RowHandler method rowToMessage.

/* Convert a row to a message. If Schema payload field is a Row type, payloadSerializer is required. */
PubSubMessage rowToMessage(Row row) {
    row = castRow(row, row.getSchema(), schema);
    PubSubMessage.Builder builder = PubSubMessage.newBuilder();
    if (schema.hasField(MESSAGE_KEY_FIELD)) {
        byte[] bytes = row.getBytes(MESSAGE_KEY_FIELD);
        if (bytes != null) {
            builder.setKey(ByteString.copyFrom(bytes));
        }
    }
    if (schema.hasField(EVENT_TIMESTAMP_FIELD)) {
        ReadableDateTime time = row.getDateTime(EVENT_TIMESTAMP_FIELD);
        if (time != null) {
            builder.setEventTime(Timestamps.fromMillis(time.getMillis()));
        }
    }
    if (schema.hasField(ATTRIBUTES_FIELD)) {
        Collection<Row> attributes = row.getArray(ATTRIBUTES_FIELD);
        if (attributes != null) {
            attributes.forEach(entry -> {
                AttributeValues.Builder valuesBuilder = AttributeValues.newBuilder();
                Collection<byte[]> values = checkArgumentNotNull(entry.getArray(ATTRIBUTES_VALUES_FIELD));
                values.forEach(bytes -> valuesBuilder.addValues(ByteString.copyFrom(bytes)));
                builder.putAttributes(checkArgumentNotNull(entry.getString(ATTRIBUTES_KEY_FIELD)), valuesBuilder.build());
            });
        }
    }
    if (payloadSerializer == null) {
        byte[] payload = row.getBytes(PAYLOAD_FIELD);
        if (payload != null) {
            builder.setData(ByteString.copyFrom(payload));
        }
    } else {
        Row payload = row.getRow(PAYLOAD_FIELD);
        if (payload != null) {
            builder.setData(ByteString.copyFrom(payloadSerializer.serialize(payload)));
        }
    }
    return builder.build();
}

Example 4 with PubSubMessage

use of com.google.cloud.pubsublite.proto.PubSubMessage in project beam by apache.

the class RowHandlerTest method fullRowToMessage.

@Test
public void fullRowToMessage() {
    RowHandler rowHandler = new RowHandler(FULL_WRITE_SCHEMA);
    Instant now = Instant.now();
    Row row = Row.withSchema(FULL_WRITE_SCHEMA).withFieldValue(RowHandler.MESSAGE_KEY_FIELD, "val1".getBytes(UTF_8)).withFieldValue(RowHandler.PAYLOAD_FIELD, "val2".getBytes(UTF_8)).withFieldValue(RowHandler.EVENT_TIMESTAMP_FIELD, now).withFieldValue(RowHandler.ATTRIBUTES_FIELD, ImmutableList.of(Row.withSchema(RowHandler.ATTRIBUTES_ENTRY_SCHEMA).attachValues("key1", ImmutableList.of("attr1".getBytes(UTF_8), "attr2".getBytes(UTF_8))), Row.withSchema(RowHandler.ATTRIBUTES_ENTRY_SCHEMA).attachValues("key2", ImmutableList.of("attr3".getBytes(UTF_8))))).build();
    PubSubMessage expected = PubSubMessage.newBuilder().setKey(ByteString.copyFromUtf8("val1")).setData(ByteString.copyFromUtf8("val2")).setEventTime(Timestamps.fromMillis(now.getMillis())).putAttributes("key1", AttributeValues.newBuilder().addValues(ByteString.copyFromUtf8("attr1")).addValues(ByteString.copyFromUtf8("attr2")).build()).putAttributes("key2", AttributeValues.newBuilder().addValues(ByteString.copyFromUtf8("attr3")).build()).build();
    assertEquals(expected, rowHandler.rowToMessage(row));
}

Example 5 with PubSubMessage

use of com.google.cloud.pubsublite.proto.PubSubMessage in project beam by apache.

the class PubsubLiteIO method write.

/**
 * Write messages to Pub/Sub Lite.
 *
 * <pre>{@code
 * TopicPath topicPath =
 *         TopicPath.newBuilder()
 *             .setProjectNumber(projectNum)
 *             .setLocation(zone)
 *             .setName(topicName)
 *             .build();
 *
 * PCollection<Message> messages = ...;
 * messages.apply(PubsubLiteIO.write(
 *     PublisherOptions.newBuilder().setTopicPath(topicPath).build());
 *
 * }</pre>
 */
public static PTransform<PCollection<PubSubMessage>, PDone> write(PublisherOptions options) {
    return new PTransform<PCollection<PubSubMessage>, PDone>() {

        @Override
        public PDone expand(PCollection<PubSubMessage> input) {
            PubsubLiteSink sink = new PubsubLiteSink(options);
            input.apply("Write", ParDo.of(sink));
            return PDone.in(input.getPipeline());
        }
    };
}