Modelo de Pub/Sub Proto para BigQuery

O modelo Pub/Sub para BigQuery é um pipeline de streaming que carrega dados proto de uma subscrição do Pub/Sub para uma tabela do BigQuery. Todos os erros que ocorrem durante a escrita na tabela do BigQuery são transmitidos para um tópico não processado do Pub/Sub.

Pode fornecer uma função definida pelo utilizador (FDU) em JavaScript para transformar dados. Os erros durante a execução da UDF podem ser enviados para um tópico do Pub/Sub separado ou para o mesmo tópico não processado que os erros do BigQuery.

Antes de executar um pipeline do Dataflow para este cenário, considere se uma subscrição do Pub/Sub BigQuery com uma UDF cumpre os seus requisitos.

Requisitos do pipeline

• A subscrição do Pub/Sub de entrada tem de existir.
• O ficheiro de esquema dos registos Proto tem de existir no Cloud Storage.
• O tópico Pub/Sub de saída tem de existir.
• O conjunto de dados do BigQuery de saída tem de existir.
• Se a tabela do BigQuery existir, tem de ter um esquema correspondente aos dados proto, independentemente do valor de createDisposition.

Parâmetros de modelos

Parâmetros obrigatórios

• protoSchemaPath (caminho do Cloud Storage para o ficheiro de esquema Proto): caminho do Cloud Storage para um ficheiro de conjunto de descritores autónomo. Exemplo: gs://MyBucket/schema.pb. schema.pb pode ser gerado adicionando --descriptor_set_out=schema.pb ao comando protoc que compila os protos. A flag --include_imports pode ser usada para garantir que o ficheiro é autónomo.
• fullMessageName (nome completo da mensagem Proto): o nome completo da mensagem (por exemplo: package.name.MessageName). Se a mensagem estiver aninhada noutra mensagem, inclua todas as mensagens com o delimitador "." (exemplo: package.name.OuterMessage.InnerMessage). "package.name" deve ser da declaração package e não da declaração java_package.
• inputSubscription (subscrição de entrada do Pub/Sub): subscrição do Pub/Sub para ler a entrada, no formato "projects/your-project-id/subscriptions/your-subscription-name" (exemplo: projects/your-project-id/subscriptions/your-subscription-name).
• outputTableSpec (tabela de resultados do BigQuery): localização da tabela do BigQuery para escrever os resultados. O nome deve estar no formato <project>:<dataset>.<table_name>. O esquema da tabela tem de corresponder aos objetos de entrada.
• outputTopic (tópico de saída do Pub/Sub): o nome do tópico para o qual os dados devem ser publicados, no formato "projects/your-project-id/topics/your-topic-name" (exemplo: projects/your-project-id/topics/your-topic-name).

Parâmetros opcionais

• preserveProtoFieldNames (Preservar nomes de campos proto): indicador para controlar se os nomes de campos proto devem ser mantidos ou convertidos para lowerCamelCase. Se a tabela já existir, isto deve basear-se no que corresponde ao esquema da tabela. Caso contrário, determina os nomes das colunas da tabela criada. Verdadeiro para preservar o snake_case do proto. False converte os campos em lowerCamelCase. (Predefinição: false).
• bigQueryTableSchemaPath (caminho do esquema da tabela do BigQuery): caminho do Cloud Storage para o ficheiro JSON do esquema do BigQuery. Se não estiver definido, o esquema é inferido a partir do esquema Proto. (Exemplo: gs://MyBucket/bq_schema.json).
• udfOutputTopic (tópico de saída do Pub/Sub para falhas de UDF): um tópico de saída opcional para enviar falhas de UDF. Se esta opção não estiver definida, as falhas são escritas no mesmo tópico que as falhas do BigQuery. (Exemplo: projects/your-project-id/topics/your-topic-name).
• writeDisposition (Write Disposition a usar para o BigQuery): WriteDisposition do BigQuery. Por exemplo, WRITE_APPEND, WRITE_EMPTY ou WRITE_TRUNCATE. A predefinição é: WRITE_APPEND.
• createDisposition (Create Disposition to use for BigQuery): BigQuery CreateDisposition. Por exemplo, CREATE_IF_NEEDED, CREATE_NEVER. A predefinição é: CREATE_IF_NEEDED.
• javascriptTextTransformGcsPath (caminho do Cloud Storage para a origem da FDU JavaScript): o padrão do caminho do Cloud Storage para o código JavaScript que contém as suas funções definidas pelo utilizador. (Exemplo: gs://your-bucket/your-function.js).
• javascriptTextTransformFunctionName (nome da função JavaScript da FDU): o nome da função a chamar a partir do seu ficheiro JavaScript. Use apenas letras, dígitos e sublinhados. (Exemplo: "transform" ou "transform_udf1").
• javascriptTextTransformReloadIntervalMinutes (intervalo de recarregamento automático da FDU de JavaScript [minutos]): define o intervalo que os trabalhadores podem verificar quanto a alterações da FDU de JavaScript para recarregar os ficheiros. A predefinição é: 0.
• useStorageWriteApi (Usar API Storage Write do BigQuery): se for verdadeiro, o pipeline usa a API Storage Write ao escrever os dados no BigQuery (consulte https://cloud.google.com/blog/products/data-analytics/streaming-data-into-bigquery-using-storage-write-api). O valor predefinido é false. Quando usar a API Storage Write no modo exatamente uma vez, tem de definir os seguintes parâmetros: "Número de streams para a API Storage Write do BigQuery" e "Frequência de acionamento em segundos para a API Storage Write do BigQuery". Se ativar o modo Dataflow at-least-once ou definir o parâmetro useStorageWriteApiAtLeastOnce como verdadeiro, não precisa de definir o número de streams nem a frequência de acionamento.
• useStorageWriteApiAtLeastOnce (Usar semântica de pelo menos uma vez na API Storage Write do BigQuery): este parâmetro só entra em vigor se a opção "Usar API Storage Write do BigQuery" estiver ativada. Se estiver ativada, a semântica de pelo menos uma vez é usada para a API Storage Write. Caso contrário, é usada a semântica de exatamente uma vez. A predefinição é: false.
• numStorageWriteApiStreams (Número de streams para a API Storage Write do BigQuery): o número de streams define o paralelismo da transformação de escrita do BigQueryIO e corresponde aproximadamente ao número de streams da API Storage Write que serão usadas pelo pipeline. Consulte https://cloud.google.com/blog/products/data-analytics/streaming-data-into-bigquery-using-storage-write-api para ver os valores recomendados. A predefinição é: 0.
• storageWriteApiTriggeringFrequencySec (Frequência de acionamento em segundos para a API BigQuery Storage Write): a frequência de acionamento determina a rapidez com que os dados ficam visíveis para consulta no BigQuery. Consulte https://cloud.google.com/blog/products/data-analytics/streaming-data-into-bigquery-using-storage-write-api para ver os valores recomendados.

Função definida pelo utilizador

Opcionalmente, pode estender este modelo escrevendo uma função definida pelo utilizador (FDU). O modelo chama a FDU para cada elemento de entrada. Os payloads dos elementos são serializados como strings JSON. Para mais informações, consulte o artigo Crie funções definidas pelo utilizador para modelos do Dataflow.

Especificação da função

A FDU tem a seguinte especificação:

Entrada: o campo de dados da mensagem do Pub/Sub, serializado como uma string JSON.

Saída: uma string JSON que corresponde ao esquema da tabela de destino do BigQuery.

Execute o modelo

Consola

1. Aceda à página do fluxo de dados Criar tarefa a partir de um modelo.
Aceda a Criar tarefa a partir de modelo
2. No campo Nome da tarefa, introduza um nome exclusivo para a tarefa.
3. Opcional: para Ponto final regional, selecione um valor no menu pendente. A região predefinida é us-central1.

   Para ver uma lista das regiões onde pode executar uma tarefa do Dataflow, consulte as localizações do Dataflow.

4. No menu pendente Modelo do fluxo de dados, selecione the Pub/Sub Proto to BigQuery template.
5. Nos campos de parâmetros fornecidos, introduza os valores dos parâmetros.
6. Clique em Executar tarefa.

gcloud

Nota: para usar a CLI do Google Cloud para executar modelos flex, tem de ter a versão CLI do Google Cloud 284.0.0 ou posterior.

Na shell ou no terminal, execute o modelo:

gcloud dataflow flex-template run JOB_NAME \
    --region=REGION_NAME \
    --template-file-gcs-location=gs://dataflow-templates-REGION_NAME/VERSION/flex/PubSub_Proto_to_BigQuery_Flex \
    --parameters \
schemaPath=SCHEMA_PATH,\
fullMessageName=PROTO_MESSAGE_NAME,\
inputSubscription=SUBSCRIPTION_NAME,\
outputTableSpec=BIGQUERY_TABLE,\
outputTopic=UNPROCESSED_TOPIC
  

Substitua o seguinte:

• JOB_NAME: um nome de tarefa exclusivo à sua escolha
• REGION_NAME: a região onde quer implementar a tarefa do Dataflow, por exemplo, us-central1
• VERSION: a versão do modelo que quer usar

  Pode usar os seguintes valores:

  • latest para usar a versão mais recente do modelo, que está disponível na pasta principal sem data no contentor: gs://dataflow-templates-REGION_NAME/latest/
  • o nome da versão, como 2023-09-12-00_RC00, para usar uma versão específica do modelo, que pode ser encontrada aninhada na pasta principal com a data correspondente no contentor: gs://dataflow-templates-REGION_NAME/
  Atenção: a versão mais recente dos modelos pode ser atualizada com alterações destrutivas. Os seus ambientes de produção devem usar modelos mantidos na pasta principal datada mais recente para evitar que estas alterações significativas afetem os seus fluxos de trabalho de produção.
• SCHEMA_PATH: o caminho do Cloud Storage para o ficheiro de esquema Proto (por exemplo, gs://MyBucket/file.pb)
• PROTO_MESSAGE_NAME: o nome da mensagem Proto (por exemplo, package.name.MessageName)
• SUBSCRIPTION_NAME: o nome da subscrição de entrada do Pub/Sub
• BIGQUERY_TABLE: o nome da tabela de saída do BigQuery
• UNPROCESSED_TOPIC: o tópico do Pub/Sub a usar para a fila não processada

API

Para executar o modelo através da API REST, envie um pedido HTTP POST. Para mais informações sobre a API e os respetivos âmbitos de autorização, consulte projects.templates.launch.

POST https://dataflow.googleapis.com/v1b3/projects/PROJECT_ID/locations/LOCATION/flexTemplates:launch
{
   "launch_parameter": {
      "jobName": "JOB_NAME",
      "containerSpecGcsPath": "gs://dataflow-templates-REGION_NAME/VERSION/flex/PubSub_Proto_to_BigQuery_Flex",
      "parameters": {
          "schemaPath": "SCHEMA_PATH",
          "fullMessageName": "PROTO_MESSAGE_NAME",
          "inputSubscription": "SUBSCRIPTION_NAME",
          "outputTableSpec": "BIGQUERY_TABLE",
          "outputTopic": "UNPROCESSED_TOPIC"
      }
   }
}
  

Substitua o seguinte:

• PROJECT_ID: o ID do projeto onde quer executar a tarefa do Dataflow Google Cloud
• JOB_NAME: um nome de tarefa exclusivo à sua escolha
• LOCATION: a região onde quer implementar a tarefa do Dataflow, por exemplo, us-central1
• VERSION: a versão do modelo que quer usar

  Pode usar os seguintes valores:

  • latest para usar a versão mais recente do modelo, que está disponível na pasta principal sem data no contentor: gs://dataflow-templates-REGION_NAME/latest/
  • o nome da versão, como 2023-09-12-00_RC00, para usar uma versão específica do modelo, que pode ser encontrada aninhada na pasta principal com a data correspondente no contentor: gs://dataflow-templates-REGION_NAME/
  Atenção: a versão mais recente dos modelos pode ser atualizada com alterações destrutivas. Os seus ambientes de produção devem usar modelos mantidos na pasta principal datada mais recente para evitar que estas alterações significativas afetem os seus fluxos de trabalho de produção.
• SCHEMA_PATH: o caminho do Cloud Storage para o ficheiro de esquema Proto (por exemplo, gs://MyBucket/file.pb)
• PROTO_MESSAGE_NAME: o nome da mensagem Proto (por exemplo, package.name.MessageName)
• SUBSCRIPTION_NAME: o nome da subscrição de entrada do Pub/Sub
• BIGQUERY_TABLE: o nome da tabela de saída do BigQuery
• UNPROCESSED_TOPIC: o tópico do Pub/Sub a usar para a fila não processada

Código-fonte do modelo

Java

/*
 * 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());
  }
}

O que se segue?

• Saiba mais sobre os modelos do Dataflow.
• Consulte a lista de modelos fornecidos pela Google.