Pub/Sub Proto to BigQuery with Python UDF template

/*  * Copyright (C) 2021 Google LLC  *  * Licensed 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 com.google.cloud.teleport.v2.templates; import static com.google.cloud.teleport.v2.transforms.PythonExternalTextTransformer.FailsafeRowPythonExternalUdf.stringMappingFunction; import static java.nio.charset.StandardCharsets.UTF_8; import com.google.cloud.teleport.metadata.MultiTemplate; import com.google.cloud.teleport.metadata.Template; import com.google.cloud.teleport.metadata.TemplateCategory; import com.google.cloud.teleport.metadata.TemplateParameter; import com.google.cloud.teleport.v2.coders.FailsafeElementCoder; import com.google.cloud.teleport.v2.common.UncaughtExceptionLogger; import com.google.cloud.teleport.v2.options.BigQueryCommonOptions.WriteOptions; import com.google.cloud.teleport.v2.options.BigQueryStorageApiStreamingOptions; import com.google.cloud.teleport.v2.options.PubsubCommonOptions.ReadSubscriptionOptions; import com.google.cloud.teleport.v2.options.PubsubCommonOptions.WriteTopicOptions; import com.google.cloud.teleport.v2.templates.PubsubProtoToBigQuery.PubSubProtoToBigQueryOptions; import com.google.cloud.teleport.v2.transforms.BigQueryConverters; import com.google.cloud.teleport.v2.transforms.ErrorConverters; import com.google.cloud.teleport.v2.transforms.FailsafeElementTransforms.ConvertFailsafeElementToPubsubMessage; import com.google.cloud.teleport.v2.transforms.JavascriptTextTransformer.FailsafeJavascriptUdf; import com.google.cloud.teleport.v2.transforms.PythonExternalTextTransformer; import com.google.cloud.teleport.v2.transforms.PythonExternalTextTransformer.RowToStringFailsafeElementFn; import com.google.cloud.teleport.v2.utils.BigQueryIOUtils; import com.google.cloud.teleport.v2.utils.GCSUtils; import com.google.cloud.teleport.v2.utils.SchemaUtils; import com.google.cloud.teleport.v2.values.FailsafeElement; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Strings; import com.google.protobuf.Descriptors.Descriptor; import com.google.protobuf.DynamicMessage; import com.google.protobuf.InvalidProtocolBufferException; import com.google.protobuf.util.JsonFormat; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.coders.NullableCoder; import org.apache.beam.sdk.coders.RowCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write; import org.apache.beam.sdk.io.gcp.bigquery.WriteResult; import org.apache.beam.sdk.io.gcp.pubsub.PubsubIO; import org.apache.beam.sdk.io.gcp.pubsub.PubsubIO.Read; import org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage; import org.apache.beam.sdk.options.Default; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.options.Validation.Required; import org.apache.beam.sdk.transforms.MapElements; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionTuple; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.TupleTagList; import org.apache.beam.sdk.values.TypeDescriptor; import org.apache.beam.sdk.values.TypeDescriptors; import org.apache.commons.lang3.ArrayUtils; /**  * A template for writing <a href="https://developers.google.com/protocol-buffers">Protobuf</a>  * records from Pub/Sub to BigQuery.  *  * <p>Persistent failures are written to a Pub/Sub unprocessed topic.  *  * <p>Check out <a  * href="https://github.com/GoogleCloudPlatform/DataflowTemplates/blob/main/v2/pubsub-binary-to-bigquery/README_PubSub_Proto_to_BigQuery.md">README</a>  * for instructions on how to use or modify this template.  */ @MultiTemplate({  @Template(  name = "PubSub_Proto_to_BigQuery_Flex",  category = TemplateCategory.STREAMING,  displayName = "Pub/Sub Proto to BigQuery",  description = {  "The Pub/Sub proto to BigQuery template is a streaming pipeline that ingests proto data from a Pub/Sub subscription into a BigQuery table. "  + "Any errors that occur while writing to the BigQuery table are streamed into a Pub/Sub unprocessed topic.\n",  "A JavaScript user-defined function (UDF) can be provided to transform data. "  + "Errors while executing the UDF can be sent to either a separate Pub/Sub topic or the same unprocessed topic as the BigQuery errors."  },  skipOptions = {  "pythonExternalTextTransformGcsPath",  "pythonExternalTextTransformFunctionName"  },  optionsClass = PubSubProtoToBigQueryOptions.class,  flexContainerName = "pubsub-proto-to-bigquery",  documentation =  "https://cloud.google.com/dataflow/docs/guides/templates/provided/pubsub-proto-to-bigquery",  contactInformation = "https://cloud.google.com/support",  requirements = {  "The input Pub/Sub subscription must exist.",  "The schema file for the Proto records must exist on Cloud Storage.",  "The output Pub/Sub topic must exist.",  "The output BigQuery dataset must exist.",  "If the BigQuery table exists, it must have a schema matching the proto data regardless of the <code>createDisposition</code> value."  },  streaming = true,  supportsAtLeastOnce = true),  @Template(  name = "PubSub_Proto_to_BigQuery_Xlang",  category = TemplateCategory.STREAMING,  displayName = "Pub/Sub Proto to BigQuery with Python UDF",  type = Template.TemplateType.XLANG,  description = {  "The Pub/Sub proto to BigQuery template is a streaming pipeline that ingests proto data from a Pub/Sub subscription into a BigQuery table. "  + "Any errors that occur while writing to the BigQuery table are streamed into a Pub/Sub unprocessed topic.\n",  "A Python user-defined function (UDF) can be provided to transform data. "  + "Errors while executing the UDF can be sent to either a separate Pub/Sub topic or the same unprocessed topic as the BigQuery errors."  },  skipOptions = {  "javascriptTextTransformGcsPath",  "javascriptTextTransformFunctionName",  "javascriptTextTransformReloadIntervalMinutes"  },  optionsClass = PubSubProtoToBigQueryOptions.class,  flexContainerName = "pubsub-proto-to-bigquery-xlang",  documentation =  "https://cloud.google.com/dataflow/docs/guides/templates/provided/pubsub-proto-to-bigquery",  contactInformation = "https://cloud.google.com/support",  requirements = {  "The input Pub/Sub subscription must exist.",  "The schema file for the Proto records must exist on Cloud Storage.",  "The output Pub/Sub topic must exist.",  "The output BigQuery dataset must exist.",  "If the BigQuery table exists, it must have a schema matching the proto data regardless of the <code>createDisposition</code> value."  },  streaming = true,  supportsAtLeastOnce = true) }) public final class PubsubProtoToBigQuery {  private static final TupleTag<FailsafeElement<String, String>> UDF_SUCCESS_TAG = new TupleTag<>();  private static final TupleTag<FailsafeElement<String, String>> UDF_FAILURE_TAG = new TupleTag<>();  private static final FailsafeElementCoder<String, String> FAILSAFE_CODER =  FailsafeElementCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of());  public static void main(String[] args) {  UncaughtExceptionLogger.register();  run(PipelineOptionsFactory.fromArgs(args).as(PubSubProtoToBigQueryOptions.class));  }  /** {@link org.apache.beam.sdk.options.PipelineOptions} for {@link PubsubProtoToBigQuery}. */  public interface PubSubProtoToBigQueryOptions  extends ReadSubscriptionOptions,  WriteOptions,  WriteTopicOptions,  PythonExternalTextTransformer.PythonExternalTextTransformerOptions,  BigQueryStorageApiStreamingOptions {  @TemplateParameter.GcsReadFile(  order = 1,  description = "Cloud Storage Path to the Proto Schema File",  helpText =  "The Cloud Storage location of the self-contained proto schema file. For example,"  + " `gs://path/to/my/file.pb`. You can generate this file with"  + " the `--descriptor_set_out` flag of the protoc command."  + " The `--include_imports` flag guarantees that the file is self-contained.")  @Required  String getProtoSchemaPath();  void setProtoSchemaPath(String value);  @TemplateParameter.Text(  order = 2,  regexes = {"^.+([a-zA-Z][a-zA-Z0-9_]+\\.?)+[a-zA-Z0-9_]$"},  description = "Full Proto Message Name",  helpText =  "The full proto message name. For example, `package.name`."  + " `MessageName`, where `package.name` is the value provided for the"  + " `package` statement and not the `java_package` statement.")  @Required  String getFullMessageName();  void setFullMessageName(String value);  @TemplateParameter.Boolean(  order = 3,  optional = true,  description = "Preserve Proto Field Names",  helpText =  "To preserve the original proto field name in JSON, set this property to `true`. "  + "To use more standard JSON names, set to `false`."  + " For example, `false` would change `field_name` to `fieldName`. Defaults to: `false`.")  @Default.Boolean(false)  Boolean getPreserveProtoFieldNames();  void setPreserveProtoFieldNames(Boolean value);  @TemplateParameter.GcsReadFile(  order = 4,  optional = true,  description = "BigQuery Table Schema Path",  helpText =  "The Cloud Storage path to the BigQuery schema path. "  + "If this value isn't provided, then the schema is inferred from the Proto schema.",  example = "gs://MyBucket/bq_schema.json")  String getBigQueryTableSchemaPath();  void setBigQueryTableSchemaPath(String value);  @TemplateParameter.PubsubTopic(  order = 5,  optional = true,  description = "Pub/Sub output topic for UDF failures",  helpText =  "The Pub/Sub topic storing the UDF errors."  + " If this value isn't provided, UDF errors are sent to the same topic as `outputTopic`.",  example = "projects/your-project-id/topics/your-topic-name")  String getUdfOutputTopic();  void setUdfOutputTopic(String udfOutputTopic);  // Hide the UseStorageWriteApiAtLeastOnce in the UI, because it will automatically be turned  // on when pipeline is running on ALO mode and using the Storage Write API  @TemplateParameter.Boolean(  order = 6,  optional = true,  parentName = "useStorageWriteApi",  parentTriggerValues = {"true"},  description = "Use at at-least-once semantics in BigQuery Storage Write API",  helpText =  "When using the Storage Write API, specifies the write semantics."  + " To use at-least-once semantics (https://beam.apache.org/documentation/io/built-in/google-bigquery/#at-least-once-semantics), set this parameter to true`. To use exactly-once semantics, set the parameter to `false`."  + " This parameter applies only when `useStorageWriteApi` is `true`. The default value is `false`.",  hiddenUi = true)  @Default.Boolean(false)  @Override  Boolean getUseStorageWriteApiAtLeastOnce();  void setUseStorageWriteApiAtLeastOnce(Boolean value);  }  /** Runs the pipeline and returns the results. */  private static PipelineResult run(PubSubProtoToBigQueryOptions options) {  BigQueryIOUtils.validateBQStorageApiOptionsStreaming(options);  Pipeline pipeline = Pipeline.create(options);  Descriptor descriptor = getDescriptor(options);  PCollection<String> maybeForUdf =  pipeline  .apply("Read From Pubsub", readPubsubMessages(options, descriptor))  .apply("Dynamic Message to TableRow", new ConvertDynamicProtoMessageToJson(options));  WriteResult writeResult =  runUdf(maybeForUdf, options)  .apply("Write to BigQuery", writeToBigQuery(options, descriptor));  BigQueryIOUtils.writeResultToBigQueryInsertErrors(writeResult, options)  .apply(  "Create Error Payload",  ErrorConverters.BigQueryInsertErrorToPubsubMessage.<String>newBuilder()  .setPayloadCoder(StringUtf8Coder.of())  .setTranslateFunction(BigQueryConverters::tableRowToJson)  .build())  .apply("Write Failed BQ Records", PubsubIO.writeMessages().to(options.getOutputTopic()));  return pipeline.run();  }  /** Gets the {@link Descriptor} for the message type in the Pub/Sub topic. */  @VisibleForTesting  static Descriptor getDescriptor(PubSubProtoToBigQueryOptions options) {  String schemaPath = options.getProtoSchemaPath();  String messageName = options.getFullMessageName();  Descriptor descriptor = SchemaUtils.getProtoDomain(schemaPath).getDescriptor(messageName);  if (descriptor == null) {  throw new IllegalArgumentException(  messageName + " is not a recognized message in " + schemaPath);  }  return descriptor;  }  /** Returns the {@link PTransform} for reading Pub/Sub messages. */  private static Read<DynamicMessage> readPubsubMessages(  PubSubProtoToBigQueryOptions options, Descriptor descriptor) {  return PubsubIO.readProtoDynamicMessages(descriptor)  .fromSubscription(options.getInputSubscription())  .withDeadLetterTopic(options.getOutputTopic());  }  /**  * Writes messages to BigQuery, creating the table if necessary and allowed in {@code options}.  *  * <p>The BigQuery schema will be inferred from {@code descriptor} unless a JSON schema path is  * specified in {@code options}.  */  @VisibleForTesting  static Write<String> writeToBigQuery(  PubSubProtoToBigQueryOptions options, Descriptor descriptor) {  Write<String> write =  BigQueryConverters.<String>createWriteTransform(options)  .withFormatFunction(BigQueryConverters::convertJsonToTableRow);  String schemaPath = options.getBigQueryTableSchemaPath();  if (Strings.isNullOrEmpty(schemaPath)) {  return write.withSchema(  SchemaUtils.createBigQuerySchema(descriptor, options.getPreserveProtoFieldNames()));  } else {  return write.withJsonSchema(GCSUtils.getGcsFileAsString(schemaPath));  }  }  /** {@link PTransform} that handles converting {@link PubsubMessage} values to JSON. */  private static class ConvertDynamicProtoMessageToJson  extends PTransform<PCollection<DynamicMessage>, PCollection<String>> {  private final boolean preserveProtoName;  private ConvertDynamicProtoMessageToJson(PubSubProtoToBigQueryOptions options) {  this.preserveProtoName = options.getPreserveProtoFieldNames();  }  @Override  public PCollection<String> expand(PCollection<DynamicMessage> input) {  return input.apply(  "Map to JSON",  MapElements.into(TypeDescriptors.strings())  .via(  message -> {  try {  JsonFormat.Printer printer = JsonFormat.printer();  return preserveProtoName  ? printer.preservingProtoFieldNames().print(message)  : printer.print(message);  } catch (InvalidProtocolBufferException e) {  throw new RuntimeException(e);  }  }));  }  }  /**  * Handles running the UDF.  *  * <p>If {@code options} is configured so as not to run the UDF, then the UDF will not be called.  *  * <p>This may add a branch to the pipeline for outputting failed UDF records to an unprocessed  * topic.  *  * @param jsonCollection {@link PCollection} of JSON strings for use as input to the UDF  * @param options the options containing info on running the UDF  * @return the {@link PCollection} of UDF output as JSON or {@code jsonCollection} if UDF not  * called  */  @VisibleForTesting  static PCollection<String> runUdf(  PCollection<String> jsonCollection, PubSubProtoToBigQueryOptions options) {  boolean useJavascriptUdf = !Strings.isNullOrEmpty(options.getJavascriptTextTransformGcsPath());  boolean usePythonUdf = !Strings.isNullOrEmpty(options.getPythonExternalTextTransformGcsPath());  // In order to avoid generating a graph that makes it look like a UDF was called when none was  // intended, simply return the input as "success" output.  if (!useJavascriptUdf && !usePythonUdf) {  return jsonCollection;  }  // For testing purposes, we need to do this check before creating the PTransform rather than  // in `expand`. Otherwise, we get a NullPointerException due to the PTransform not returning  // a value.  if (useJavascriptUdf  && Strings.isNullOrEmpty(options.getJavascriptTextTransformFunctionName())) {  throw new IllegalArgumentException(  "JavaScript function name cannot be null or empty if file is set");  }  if (usePythonUdf  && Strings.isNullOrEmpty(options.getPythonExternalTextTransformFunctionName())) {  throw new IllegalArgumentException(  "Python function name cannot be null or empty if file is set");  }  if (usePythonUdf && useJavascriptUdf) {  throw new IllegalArgumentException(  "Either javascript or Python gcs path must be provided, but not both.");  }  PCollectionTuple maybeSuccess;  if (usePythonUdf) {  maybeSuccess = jsonCollection.apply("Run UDF", new RunPythonUdf(options));  } else {  maybeSuccess = jsonCollection.apply("Run UDF", new RunUdf(options));  }  maybeSuccess  .get(UDF_FAILURE_TAG)  .setCoder(FAILSAFE_CODER)  .apply(  "Get UDF Failures",  ConvertFailsafeElementToPubsubMessage.<String, String>builder()  .setOriginalPayloadSerializeFn(s -> ArrayUtils.toObject(s.getBytes(UTF_8)))  .setErrorMessageAttributeKey("udfErrorMessage")  .build())  .apply("Write Failed UDF", writeUdfFailures(options));  return maybeSuccess  .get(UDF_SUCCESS_TAG)  .setCoder(FAILSAFE_CODER)  .apply(  "Get UDF Output",  MapElements.into(TypeDescriptors.strings()).via(FailsafeElement::getPayload))  .setCoder(NullableCoder.of(StringUtf8Coder.of()));  }  /** {@link PTransform} that calls a UDF and returns both success and failure output. */  private static class RunUdf extends PTransform<PCollection<String>, PCollectionTuple> {  private final PubSubProtoToBigQueryOptions options;  RunUdf(PubSubProtoToBigQueryOptions options) {  this.options = options;  }  @Override  public PCollectionTuple expand(PCollection<String> input) {  return input  .apply("Prepare Failsafe UDF", makeFailsafe())  .setCoder(FAILSAFE_CODER)  .apply(  "Call UDF",  FailsafeJavascriptUdf.<String>newBuilder()  .setFileSystemPath(options.getJavascriptTextTransformGcsPath())  .setFunctionName(options.getJavascriptTextTransformFunctionName())  .setReloadIntervalMinutes(  options.getJavascriptTextTransformReloadIntervalMinutes())  .setSuccessTag(UDF_SUCCESS_TAG)  .setFailureTag(UDF_FAILURE_TAG)  .build());  }  private static MapElements<String, FailsafeElement<String, String>> makeFailsafe() {  return MapElements.into(new TypeDescriptor<FailsafeElement<String, String>>() {})  .via((String json) -> FailsafeElement.of(json, json));  }  }  /** {@link PTransform} that calls a python UDF and returns both success and failure output. */  private static class RunPythonUdf extends PTransform<PCollection<String>, PCollectionTuple> {  private final PubSubProtoToBigQueryOptions options;  RunPythonUdf(PubSubProtoToBigQueryOptions options) {  this.options = options;  }  @Override  public PCollectionTuple expand(PCollection<String> input) {  return input  .apply("Prepare Failsafe row", stringMappingFunction())  .setCoder(  RowCoder.of(PythonExternalTextTransformer.FailsafeRowPythonExternalUdf.ROW_SCHEMA))  .apply(  "InvokeUDF",  PythonExternalTextTransformer.FailsafePythonExternalUdf.newBuilder()  .setFileSystemPath(options.getPythonExternalTextTransformGcsPath())  .setFunctionName(options.getPythonExternalTextTransformFunctionName())  .build())  .apply(  "MapRowsToFailsafeElements",  ParDo.of(new RowToStringFailsafeElementFn(UDF_SUCCESS_TAG, UDF_FAILURE_TAG))  .withOutputTags(UDF_SUCCESS_TAG, TupleTagList.of(UDF_FAILURE_TAG)));  }  }  /**  * Returns a {@link PubsubIO.Write} configured to write UDF failures to the appropriate output  * topic.  */  private static PubsubIO.Write<PubsubMessage> writeUdfFailures(  PubSubProtoToBigQueryOptions options) {  PubsubIO.Write<PubsubMessage> write = PubsubIO.writeMessages();  return Strings.isNullOrEmpty(options.getUdfOutputTopic())  ? write.to(options.getOutputTopic())  : write.to(options.getUdfOutputTopic());  } }