| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| package org.apache.beam.sdk.io.gcp.pubsub; |
| |
| import static java.util.stream.Collectors.toList; |
| import static org.apache.beam.sdk.io.gcp.pubsub.PubsubMessageToRow.ATTRIBUTES_FIELD; |
| import static org.apache.beam.sdk.io.gcp.pubsub.PubsubMessageToRow.DLQ_TAG; |
| import static org.apache.beam.sdk.io.gcp.pubsub.PubsubMessageToRow.MAIN_TAG; |
| import static org.apache.beam.sdk.io.gcp.pubsub.PubsubMessageToRow.PAYLOAD_FIELD; |
| import static org.apache.beam.sdk.io.gcp.pubsub.PubsubMessageToRow.TIMESTAMP_FIELD; |
| import static org.apache.beam.sdk.schemas.Schema.TypeName.ROW; |
| |
| import com.google.auto.service.AutoService; |
| import com.google.auto.value.AutoValue; |
| import java.io.Serializable; |
| import java.util.List; |
| import org.apache.beam.sdk.annotations.Internal; |
| import org.apache.beam.sdk.schemas.AutoValueSchema; |
| import org.apache.beam.sdk.schemas.Schema; |
| import org.apache.beam.sdk.schemas.Schema.Field; |
| import org.apache.beam.sdk.schemas.Schema.FieldType; |
| import org.apache.beam.sdk.schemas.io.InvalidConfigurationException; |
| import org.apache.beam.sdk.schemas.io.InvalidSchemaException; |
| import org.apache.beam.sdk.schemas.io.SchemaIO; |
| import org.apache.beam.sdk.schemas.io.SchemaIOProvider; |
| import org.apache.beam.sdk.schemas.io.payloads.PayloadSerializer; |
| import org.apache.beam.sdk.schemas.io.payloads.PayloadSerializers; |
| import org.apache.beam.sdk.transforms.MapElements; |
| import org.apache.beam.sdk.transforms.PTransform; |
| import org.apache.beam.sdk.values.PBegin; |
| import org.apache.beam.sdk.values.PCollection; |
| import org.apache.beam.sdk.values.PCollectionTuple; |
| import org.apache.beam.sdk.values.POutput; |
| import org.apache.beam.sdk.values.Row; |
| import org.apache.beam.sdk.values.TypeDescriptor; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap; |
| import org.checkerframework.checker.nullness.qual.Nullable; |
| |
| /** |
| * An implementation of {@link SchemaIOProvider} for reading and writing JSON/AVRO payloads with |
| * {@link PubsubIO}. |
| * |
| * <h2>Schema</h2> |
| * |
| * <p>The data schema passed to {@link #from(String, Row, Schema)} must either be of the nested or |
| * flat style. |
| * |
| * <h3>Nested style</h3> |
| * |
| * <p>If nested structure is used, the required fields included in the Pubsub message model are |
| * 'event_timestamp', 'attributes', and 'payload'. |
| * |
| * <h3>Flat style</h3> |
| * |
| * <p>If flat structure is used, the required fields include just 'event_timestamp'. Every other |
| * field is assumed part of the payload. See {@link PubsubMessageToRow} for details. |
| * |
| * <h2>Configuration</h2> |
| * |
| * <p>{@link #configurationSchema()} consists of two attributes, timestampAttributeKey and |
| * deadLetterQueue. |
| * |
| * <h3>timestampAttributeKey</h3> |
| * |
| * <p>An optional attribute key of the Pubsub message from which to extract the event timestamp. If |
| * not specified, the message publish time will be used as event timestamp. |
| * |
| * <p>This attribute has to conform to the same requirements as in {@link |
| * PubsubIO.Read.Builder#withTimestampAttribute(String)} |
| * |
| * <p>Short version: it has to be either millis since epoch or string in RFC 3339 format. |
| * |
| * <p>If the attribute is specified then event timestamps will be extracted from the specified |
| * attribute. If it is not specified then message publish timestamp will be used. |
| * |
| * <h3>deadLetterQueue</h3> |
| * |
| * <p>deadLetterQueue is an optional topic path which will be used as a dead letter queue. |
| * |
| * <p>Messages that cannot be processed will be sent to this topic. If it is not specified then |
| * exception will be thrown for errors during processing causing the pipeline to crash. |
| */ |
| @Internal |
| @AutoService(SchemaIOProvider.class) |
| @SuppressWarnings({ |
| "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402) |
| }) |
| public class PubsubSchemaIOProvider implements SchemaIOProvider { |
| public static final FieldType ATTRIBUTE_MAP_FIELD_TYPE = |
| Schema.FieldType.map(FieldType.STRING.withNullable(false), FieldType.STRING); |
| public static final Schema ATTRIBUTE_ARRAY_ENTRY_SCHEMA = |
| Schema.builder().addStringField("key").addStringField("value").build(); |
| public static final FieldType ATTRIBUTE_ARRAY_FIELD_TYPE = |
| Schema.FieldType.array(Schema.FieldType.row(ATTRIBUTE_ARRAY_ENTRY_SCHEMA)); |
| |
| /** Returns an id that uniquely represents this IO. */ |
| @Override |
| public String identifier() { |
| return "pubsub"; |
| } |
| |
| /** |
| * Returns the expected schema of the configuration object. Note this is distinct from the schema |
| * of the data source itself. |
| */ |
| @Override |
| public Schema configurationSchema() { |
| return Schema.builder() |
| .addNullableField("timestampAttributeKey", FieldType.STRING) |
| .addNullableField("deadLetterQueue", FieldType.STRING) |
| .addNullableField("format", FieldType.STRING) |
| // For ThriftPayloadSerializerProvider |
| .addNullableField("thriftClass", FieldType.STRING) |
| .addNullableField("thriftProtocolFactoryClass", FieldType.STRING) |
| // For ProtoPayloadSerializerProvider |
| .addNullableField("protoClass", FieldType.STRING) |
| .build(); |
| } |
| |
| /** |
| * Produce a SchemaIO given a String representing the data's location, the schema of the data that |
| * resides there, and some IO-specific configuration object. |
| */ |
| @Override |
| public PubsubSchemaIO from(String location, Row configuration, Schema dataSchema) { |
| validateConfigurationSchema(configuration); |
| validateDlq(configuration.getValue("deadLetterQueue")); |
| validateDataSchema(dataSchema); |
| return new PubsubSchemaIO(location, configuration, dataSchema); |
| } |
| |
| @Override |
| public boolean requiresDataSchema() { |
| return true; |
| } |
| |
| @Override |
| public PCollection.IsBounded isBounded() { |
| return PCollection.IsBounded.UNBOUNDED; |
| } |
| |
| private void validateDataSchema(Schema schema) { |
| if (schema == null) { |
| throw new InvalidSchemaException( |
| "Unsupported schema specified for Pubsub source in CREATE TABLE." |
| + "CREATE TABLE for Pubsub topic must not be null"); |
| } |
| if (!PubsubSchemaIO.fieldPresent(schema, TIMESTAMP_FIELD, FieldType.DATETIME)) { |
| throw new InvalidSchemaException( |
| "Unsupported schema specified for Pubsub source in CREATE TABLE." |
| + "CREATE TABLE for Pubsub topic must include at least 'event_timestamp' field of " |
| + "type 'TIMESTAMP'"); |
| } |
| } |
| |
| private void validateDlq(String deadLetterQueue) { |
| if (deadLetterQueue != null && deadLetterQueue.isEmpty()) { |
| throw new InvalidConfigurationException("Dead letter queue topic name is not specified"); |
| } |
| } |
| |
| private void validateConfigurationSchema(Row configuration) { |
| if (!configuration.getSchema().equals(configurationSchema())) { |
| throw new InvalidConfigurationException( |
| "Configuration schema provided does not match expected"); |
| } |
| } |
| |
| /** An abstraction to create schema aware IOs. */ |
| private static class PubsubSchemaIO implements SchemaIO, Serializable { |
| protected final Schema dataSchema; |
| protected final String location; |
| protected final boolean useFlatSchema; |
| protected final Config config; |
| |
| private PubsubSchemaIO(String location, Row config, Schema dataSchema) { |
| this.dataSchema = dataSchema; |
| this.location = location; |
| this.useFlatSchema = !shouldUseNestedSchema(dataSchema); |
| this.config = |
| new AutoValueSchema().fromRowFunction(TypeDescriptor.of(Config.class)).apply(config); |
| } |
| |
| @Override |
| public Schema schema() { |
| return dataSchema; |
| } |
| |
| private boolean needsSerializer() { |
| return useFlatSchema || !fieldPresent(schema(), PAYLOAD_FIELD, FieldType.BYTES); |
| } |
| |
| @Override |
| public PTransform<PBegin, PCollection<Row>> buildReader() { |
| return new PTransform<PBegin, PCollection<Row>>() { |
| @Override |
| public PCollection<Row> expand(PBegin begin) { |
| PubsubMessageToRow.Builder builder = |
| PubsubMessageToRow.builder() |
| .messageSchema(dataSchema) |
| .useDlq(config.useDeadLetterQueue()) |
| .useFlatSchema(useFlatSchema); |
| if (needsSerializer()) { |
| builder.serializerProvider(config::serializer); |
| } |
| PCollectionTuple rowsWithDlq = |
| begin |
| .apply("ReadFromPubsub", readMessagesWithAttributes()) |
| .apply("PubsubMessageToRow", builder.build()); |
| rowsWithDlq.get(MAIN_TAG).setRowSchema(dataSchema); |
| |
| if (config.useDeadLetterQueue()) { |
| rowsWithDlq.get(DLQ_TAG).apply(writeMessagesToDlq()); |
| } |
| |
| return rowsWithDlq.get(MAIN_TAG); |
| } |
| }; |
| } |
| |
| @Override |
| public PTransform<PCollection<Row>, POutput> buildWriter() { |
| @Nullable |
| PayloadSerializer serializer = |
| needsSerializer() ? config.serializer(stripFromTimestampField(dataSchema)) : null; |
| |
| return new PTransform<PCollection<Row>, POutput>() { |
| @Override |
| public POutput expand(PCollection<Row> input) { |
| PCollection<Row> filtered = |
| input.apply(new AddTimestampAttribute(config.useTimestampAttribute())); |
| PCollection<PubsubMessage> transformed; |
| if (useFlatSchema) { |
| transformed = |
| filtered.apply( |
| "Transform Flat Schema", |
| MapElements.into(TypeDescriptor.of(PubsubMessage.class)) |
| .via( |
| row -> |
| new PubsubMessage(serializer.serialize(row), ImmutableMap.of()))); |
| } else { |
| transformed = |
| filtered.apply( |
| "Transform Nested Schema", |
| MapElements.via(new NestedRowToMessage(serializer, filtered.getSchema()))); |
| } |
| return transformed.apply(createPubsubMessageWrite()); |
| } |
| }; |
| } |
| |
| private PubsubIO.Read<PubsubMessage> readMessagesWithAttributes() { |
| PubsubIO.Read<PubsubMessage> read = PubsubIO.readMessagesWithAttributes().fromTopic(location); |
| |
| return config.useTimestampAttribute() |
| ? read.withTimestampAttribute(config.getTimestampAttributeKey()) |
| : read; |
| } |
| |
| private PubsubIO.Write<PubsubMessage> createPubsubMessageWrite() { |
| PubsubIO.Write<PubsubMessage> write = PubsubIO.writeMessages().to(location); |
| if (config.useTimestampAttribute()) { |
| write = write.withTimestampAttribute(config.getTimestampAttributeKey()); |
| } |
| return write; |
| } |
| |
| private PubsubIO.Write<PubsubMessage> writeMessagesToDlq() { |
| PubsubIO.Write<PubsubMessage> write = |
| PubsubIO.writeMessages().to(config.getDeadLetterQueue()); |
| |
| return config.useTimestampAttribute() |
| ? write.withTimestampAttribute(config.getTimestampAttributeKey()) |
| : write; |
| } |
| |
| private boolean hasValidAttributesField(Schema schema) { |
| return fieldPresent(schema, ATTRIBUTES_FIELD, ATTRIBUTE_MAP_FIELD_TYPE) |
| || fieldPresent(schema, ATTRIBUTES_FIELD, ATTRIBUTE_ARRAY_FIELD_TYPE); |
| } |
| |
| private boolean hasValidPayloadField(Schema schema) { |
| if (!schema.hasField(PAYLOAD_FIELD)) { |
| return false; |
| } |
| if (fieldPresent(schema, PAYLOAD_FIELD, FieldType.BYTES)) { |
| return true; |
| } |
| return schema.getField(PAYLOAD_FIELD).getType().getTypeName().equals(ROW); |
| } |
| |
| private boolean shouldUseNestedSchema(Schema schema) { |
| return hasValidPayloadField(schema) && hasValidAttributesField(schema); |
| } |
| |
| private static boolean fieldPresent( |
| Schema schema, String field, Schema.FieldType expectedType) { |
| return schema.hasField(field) |
| && expectedType.equivalent( |
| schema.getField(field).getType(), Schema.EquivalenceNullablePolicy.IGNORE); |
| } |
| } |
| |
| private static Schema stripFromTimestampField(Schema schema) { |
| List<Field> selectedFields = |
| schema.getFields().stream() |
| .filter(field -> !TIMESTAMP_FIELD.equals(field.getName())) |
| .collect(toList()); |
| return Schema.of(selectedFields.toArray(new Schema.Field[0])); |
| } |
| |
| @AutoValue |
| abstract static class Config implements Serializable { |
| |
| abstract @Nullable String getTimestampAttributeKey(); |
| |
| abstract @Nullable String getDeadLetterQueue(); |
| |
| abstract @Nullable String getFormat(); |
| |
| // For ThriftPayloadSerializerProvider |
| abstract @Nullable String getThriftClass(); |
| |
| abstract @Nullable String getThriftProtocolFactoryClass(); |
| |
| // For ProtoPayloadSerializerProvider |
| abstract @Nullable String getProtoClass(); |
| |
| boolean useDeadLetterQueue() { |
| return getDeadLetterQueue() != null; |
| } |
| |
| boolean useTimestampAttribute() { |
| return getTimestampAttributeKey() != null; |
| } |
| |
| PayloadSerializer serializer(Schema schema) { |
| String format = getFormat() == null ? "json" : getFormat(); |
| ImmutableMap.Builder<String, Object> params = ImmutableMap.builder(); |
| if (getThriftClass() != null) { |
| params.put("thriftClass", getThriftClass()); |
| } |
| if (getThriftProtocolFactoryClass() != null) { |
| params.put("thriftProtocolFactoryClass", getThriftProtocolFactoryClass()); |
| } |
| if (getProtoClass() != null) { |
| params.put("protoClass", getProtoClass()); |
| } |
| return PayloadSerializers.getSerializer(format, schema, params.build()); |
| } |
| } |
| } |
