/* * Copyright (C) 2018 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. */packagecom.google.cloud.teleport.v2.templates;import staticcom.google.cloud.teleport.v2.transforms.BigQueryConverters.wrapBigQueryInsertError;import staticcom.google.cloud.teleport.v2.utils.GCSUtils.getGcsFileAsString;import staticcom.google.cloud.teleport.v2.utils.KMSUtils.maybeDecrypt;importcom.google.api.services.bigquery.model.ErrorProto;importcom.google.api.services.bigquery.model.TableRow;importcom.google.cloud.teleport.metadata.Template;importcom.google.cloud.teleport.metadata.TemplateCategory;importcom.google.cloud.teleport.v2.coders.FailsafeElementCoder;importcom.google.cloud.teleport.v2.common.UncaughtExceptionLogger;importcom.google.cloud.teleport.v2.options.JdbcToBigQueryOptions;importcom.google.cloud.teleport.v2.transforms.ErrorConverters;importcom.google.cloud.teleport.v2.utils.BigQueryIOUtils;importcom.google.cloud.teleport.v2.utils.GCSAwareValueProvider;importcom.google.cloud.teleport.v2.utils.JdbcConverters;importcom.google.cloud.teleport.v2.utils.ResourceUtils;importcom.google.cloud.teleport.v2.utils.SecretManagerUtils;importcom.google.cloud.teleport.v2.values.FailsafeElement;importcom.google.common.annotations.VisibleForTesting;importcom.google.common.base.Strings;importjava.util.List;importjava.util.Objects;importjava.util.stream.Collectors;importorg.apache.beam.sdk.Pipeline;importorg.apache.beam.sdk.PipelineResult;importorg.apache.beam.sdk.coders.StringUtf8Coder;importorg.apache.beam.sdk.io.FileSystems;importorg.apache.beam.sdk.io.gcp.bigquery.BigQueryIO;importorg.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write;importorg.apache.beam.sdk.io.gcp.bigquery.BigQueryInsertError;importorg.apache.beam.sdk.io.gcp.bigquery.InsertRetryPolicy;importorg.apache.beam.sdk.io.gcp.bigquery.TableRowJsonCoder;importorg.apache.beam.sdk.io.gcp.bigquery.WriteResult;importorg.apache.beam.sdk.io.jdbc.JdbcIO;importorg.apache.beam.sdk.options.PipelineOptionsFactory;importorg.apache.beam.sdk.options.ValueProvider.StaticValueProvider;importorg.apache.beam.sdk.transforms.DoFn;importorg.apache.beam.sdk.transforms.MapElements;importorg.apache.beam.sdk.transforms.ParDo;importorg.apache.beam.sdk.values.PCollection;importorg.apache.beam.sdk.values.TypeDescriptor;importorg.apache.beam.sdk.values.TypeDescriptors;importorg.joda.time.DateTime;importorg.joda.time.format.DateTimeFormat;importorg.joda.time.format.DateTimeFormatter;/** * A template that copies data from a relational database using JDBC to an existing BigQuery table. * * <p>Check out <a * href="https://github.com/GoogleCloudPlatform/DataflowTemplates/blob/main/v2/jdbc-and-googlecloud/README_Jdbc_to_BigQuery_Flex.md">README</a> * for instructions on how to use or modify this template. */@Template(name="Jdbc_to_BigQuery_Flex",category=TemplateCategory.BATCH,displayName="JDBC to BigQuery with BigQuery Storage API support",description={"The JDBC to BigQuery template is a batch pipeline that copies data from a relational database table into an existing BigQuery table. "+"This pipeline uses JDBC to connect to the relational database. You can use this template to copy data from any relational database with available JDBC drivers into BigQuery.","For an extra layer of protection, you can also pass in a Cloud KMS key along with a Base64-encoded username, password, and connection string parameters encrypted with the Cloud KMS key. "+"See the <a href=\"https://cloud.google.com/kms/docs/reference/rest/v1/projects.locations.keyRings.cryptoKeys/encrypt\">Cloud KMS API encryption endpoint</a> for additional details on encrypting your username, password, and connection string parameters."},optionsClass=JdbcToBigQueryOptions.class,flexContainerName="jdbc-and-googlecloud",documentation="https://cloud.google.com/dataflow/docs/guides/templates/provided/jdbc-to-bigquery",contactInformation="https://cloud.google.com/support",preview=true,requirements={"The JDBC drivers for the relational database must be available.","If BigQuery table already exist before pipeline execution, it must have a compatible schema.","The relational database must be accessible from the subnet where Dataflow runs."})publicclassJdbcToBigQuery{/** Coder for FailsafeElement. */privatestaticfinalFailsafeElementCoder<String,String>FAILSAFE_ELEMENT_CODER=FailsafeElementCoder.of(StringUtf8Coder.of(),StringUtf8Coder.of());/**   * Main entry point for executing the pipeline. This will run the pipeline asynchronously. If   * blocking execution is required, use the {@link JdbcToBigQuery#run} method to start the pipeline   * and invoke {@code result.waitUntilFinish()} on the {@link PipelineResult}.   *   * @param args The command-line arguments to the pipeline.   */publicstaticvoidmain(String[]args){UncaughtExceptionLogger.register();// Parse the user options passed from the command-lineJdbcToBigQueryOptionsoptions=PipelineOptionsFactory.fromArgs(args).withValidation().as(JdbcToBigQueryOptions.class);run(options,writeToBQTransform(options));}/**   * Create the pipeline with the supplied options.   *   * @param options The execution parameters to the pipeline.   * @param writeToBQ the transform that outputs {@link TableRow}s to BigQuery.   * @return The result of the pipeline execution.   */@VisibleForTestingstaticPipelineResultrun(JdbcToBigQueryOptionsoptions,Write<TableRow>writeToBQ){// Validate BQ STORAGE_WRITE_API optionsBigQueryIOUtils.validateBQStorageApiOptionsBatch(options);if(!options.getUseStorageWriteApi()        &&!options.getUseStorageWriteApiAtLeastOnce()        &&!Strings.isNullOrEmpty(options.getOutputDeadletterTable())){thrownewIllegalArgumentException("outputDeadletterTable can only be specified if BigQuery Storage Write API is enabled either with useStorageWriteApi or useStorageWriteApiAtLeastOnce.");}// Create the pipelinePipelinepipeline=Pipeline.create(options);/*     * Steps: 1) Read records via JDBC and convert to TableRow via RowMapper     *        2) Append TableRow to BigQuery via BigQueryIO     */JdbcIO.DataSourceConfigurationdataSourceConfiguration=JdbcIO.DataSourceConfiguration.create(StaticValueProvider.of(options.getDriverClassName()),maybeDecrypt(maybeParseSecret(options.getConnectionURL()),options.getKMSEncryptionKey())).withUsername(maybeDecrypt(maybeParseSecret(options.getUsername()),options.getKMSEncryptionKey())).withPassword(maybeDecrypt(maybeParseSecret(options.getPassword()),options.getKMSEncryptionKey()));if(options.getDriverJars()!=null){dataSourceConfiguration=dataSourceConfiguration.withDriverJars(options.getDriverJars());}if(options.getConnectionProperties()!=null){dataSourceConfiguration=dataSourceConfiguration.withConnectionProperties(options.getConnectionProperties());}/*     * Step 1: Read records via JDBC and convert to TableRow     *         via {@link org.apache.beam.sdk.io.jdbc.JdbcIO.RowMapper}     */PCollection<TableRow>rows;if(options.getPartitionColumn()!=null &&options.getTable()!=null){// Read with PartitionsJdbcIO.ReadWithPartitions<TableRow,?>readIO=null;finalStringpartitionColumnType=options.getPartitionColumnType();if(partitionColumnType==null||"long".equals(partitionColumnType)){JdbcIO.ReadWithPartitions<TableRow,Long>longTypeReadIO=JdbcIO.<TableRow,Long>readWithPartitions(TypeDescriptors.longs()).withDataSourceConfiguration(dataSourceConfiguration).withTable(options.getTable()).withPartitionColumn(options.getPartitionColumn()).withRowMapper(JdbcConverters.getResultSetToTableRow(options.getUseColumnAlias()));if(options.getLowerBound()!=null &&options.getUpperBound()!=null){// Check if lower bound and upper bound are long type.try{longTypeReadIO=longTypeReadIO.withLowerBound(Long.valueOf(options.getLowerBound())).withUpperBound(Long.valueOf(options.getUpperBound()));}catch(NumberFormatExceptione){thrownewNumberFormatException("Expected Long values for lowerBound and upperBound, received : "+e.getMessage());}}readIO=longTypeReadIO;}elseif("datetime".equals(partitionColumnType)){JdbcIO.ReadWithPartitions<TableRow,DateTime>dateTimeReadIO=JdbcIO.<TableRow,DateTime>readWithPartitions(TypeDescriptor.of(DateTime.class)).withDataSourceConfiguration(dataSourceConfiguration).withTable(options.getTable()).withPartitionColumn(options.getPartitionColumn()).withRowMapper(JdbcConverters.getResultSetToTableRow(options.getUseColumnAlias()));if(options.getLowerBound()!=null &&options.getUpperBound()!=null){DateTimeFormatterdateFormatter=DateTimeFormat.forPattern("yyyy-MM-dd HH:mm:ss.SSSZ").withOffsetParsed();// Check if lowerBound and upperBound are DateTime type.try{dateTimeReadIO=dateTimeReadIO.withLowerBound(dateFormatter.parseDateTime(options.getLowerBound())).withUpperBound(dateFormatter.parseDateTime(options.getUpperBound()));}catch(IllegalArgumentExceptione){thrownewIllegalArgumentException("Expected DateTime values in the format for lowerBound and upperBound, received : "+e.getMessage());}}readIO=dateTimeReadIO;}else{thrownewIllegalStateException("Received unsupported partitionColumnType.");}if(options.getNumPartitions()!=null){readIO=readIO.withNumPartitions(options.getNumPartitions());}if(options.getFetchSize()!=null &&options.getFetchSize() >0){readIO=readIO.withFetchSize(options.getFetchSize());}rows=pipeline.apply("Read from JDBC with Partitions",readIO);}else{if(options.getQuery()==null){thrownewIllegalArgumentException("Either 'query' or both 'table' AND 'PartitionColumn' must be specified to read from JDBC");}JdbcIO.Read<TableRow>readIO=JdbcIO.<TableRow>read().withDataSourceConfiguration(dataSourceConfiguration).withQuery(newGCSAwareValueProvider(options.getQuery())).withCoder(TableRowJsonCoder.of()).withRowMapper(JdbcConverters.getResultSetToTableRow(options.getUseColumnAlias()));if(options.getFetchSize()!=null &&options.getFetchSize() >0){readIO=readIO.withFetchSize(options.getFetchSize());}rows=pipeline.apply("Read from JdbcIO",readIO);}/*     * Step 2: Append TableRow to an existing BigQuery table     */WriteResultwriteResult=rows.apply("Write to BigQuery",writeToBQ);/*     * Step 3.     * If using Storage Write API, capture failed inserts and either     *   a) write error rows to DLQ     *   b) fail the pipeline     */if(options.getUseStorageWriteApi()||options.getUseStorageWriteApiAtLeastOnce()){PCollection<BigQueryInsertError>insertErrors=BigQueryIOUtils.writeResultToBigQueryInsertErrors(writeResult,options);if(!Strings.isNullOrEmpty(options.getOutputDeadletterTable())){/*         * Step 3a.         * Elements that failed inserts into BigQuery are extracted and converted to FailsafeElement         */PCollection<FailsafeElement<String,String>>failedInserts=insertErrors.apply("WrapInsertionErrors",MapElements.into(FAILSAFE_ELEMENT_CODER.getEncodedTypeDescriptor()).via((BigQueryInsertErrore)->wrapBigQueryInsertError(e))).setCoder(FAILSAFE_ELEMENT_CODER);/*         * Step 3a Contd.         * Insert records that failed insert into deadletter table         */failedInserts.apply("WriteFailedRecords",ErrorConverters.WriteStringMessageErrors.newBuilder().setErrorRecordsTable(options.getOutputDeadletterTable()).setErrorRecordsTableSchema(ResourceUtils.getDeadletterTableSchemaJson()).setUseWindowedTimestamp(false).build());}else{/*         * Step 3b.         * Fail pipeline upon write errors if no DLQ was specified         */insertErrors.apply(ParDo.of(newThrowWriteErrorsDoFn()));}}// Execute the pipeline and return the result.returnpipeline.run();}staticclassThrowWriteErrorsDoFnextendsDoFn<BigQueryInsertError,Void>{@ProcessElementpublicvoidprocessElement(ProcessContextc){BigQueryInsertErrorinsertError=Objects.requireNonNull(c.element());List<String>errorMessages=insertError.getError().getErrors().stream().map(ErrorProto::getMessage).collect(Collectors.toList());StringstackTrace=String.join("\nCaused by:",errorMessages);thrownewIllegalStateException(String.format("Failed to insert row %s.\nCaused by: %s",insertError.getRow(),stackTrace));}}/**   * Create the {@link Write} transform that outputs the collection to BigQuery as per input option.   */@VisibleForTestingstaticWrite<TableRow>writeToBQTransform(JdbcToBigQueryOptionsoptions){// Needed for loading GCS filesystem before Pipeline.Create callFileSystems.setDefaultPipelineOptions(options);Write<TableRow>write=BigQueryIO.writeTableRows().withoutValidation().withCreateDisposition(Write.CreateDisposition.valueOf(options.getCreateDisposition())).withWriteDisposition(options.getIsTruncate()?Write.WriteDisposition.WRITE_TRUNCATE:Write.WriteDisposition.WRITE_APPEND).withCustomGcsTempLocation(StaticValueProvider.of(options.getBigQueryLoadingTemporaryDirectory())).withExtendedErrorInfo().to(options.getOutputTable());if(Write.CreateDisposition.valueOf(options.getCreateDisposition())!=Write.CreateDisposition.CREATE_NEVER){write=write.withJsonSchema(getGcsFileAsString(options.getBigQuerySchemaPath()));}if(options.getUseStorageWriteApi()||options.getUseStorageWriteApiAtLeastOnce()){write=write.withFailedInsertRetryPolicy(InsertRetryPolicy.retryTransientErrors());}returnwrite;}/**   * Retrieves a secret value from SecretManagerUtils if the input string matches the specified   * pattern.   *   * @param secret The input string representing a potential secret.   * @return The secret value if the input matches the pattern and the secret is found, otherwise   *     the original input string.   */privatestaticStringmaybeParseSecret(Stringsecret){// Check if the input string is not null.if(secret!=null){// Check if the input string matches the pattern for secrets stored in SecretManagerUtils.if(secret.matches("projects/.*/secrets/.*/versions/.*")){// Use .* to match any characters// Retrieve the secret value from SecretManagerUtils.returnSecretManagerUtils.getSecret(secret);}}// If the input is null or doesn't match the pattern, return the original input.returnsecret;}}