AstraDB to BigQuery template

  /* 
 * Copyright (C) 2023 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.astradb.templates 
 ; 
 import 
  
 com.datastax.oss.driver.api.core.CqlSession 
 ; 
 import 
  
 com.dtsx.astra.sdk.db.DatabaseClient 
 ; 
 import 
  
 com.google.api.services.bigquery.model. TableReference 
 
 ; 
 import 
  
 com.google.api.services.bigquery.model. TableSchema 
 
 ; 
 import 
  
 com.google.cloud.bigquery. BigQuery 
 
 ; 
 import 
  
 com.google.cloud.bigquery. BigQueryOptions 
 
 ; 
 import 
  
 com.google.cloud.bigquery. DatasetId 
 
 ; 
 import 
  
 com.google.cloud.bigquery. DatasetInfo 
 
 ; 
 import 
  
 com.google.cloud.teleport.metadata.Template 
 ; 
 import 
  
 com.google.cloud.teleport.metadata.TemplateCategory 
 ; 
 import 
  
 com.google.cloud.teleport.v2.astradb.options.AstraDbToBigQueryOptions 
 ; 
 import 
  
 com.google.cloud.teleport.v2.astradb.transforms.AstraDbToBigQueryMappingFn 
 ; 
 import 
  
 com.google.cloud.teleport.v2.common.UncaughtExceptionLogger 
 ; 
 import 
  
 com.google.cloud.teleport.v2.utils.SecretManagerUtils 
 ; 
 import 
  
 java.util.AbstractMap 
 ; 
 import 
  
 org.apache.beam.sdk.Pipeline 
 ; 
 import 
  
 org.apache.beam.sdk.coders.SerializableCoder 
 ; 
 import 
  
 org.apache.beam.sdk.io.astra.db.AstraDbIO 
 ; 
 import 
  
 org.apache.beam.sdk.io.astra.db.CqlSessionHolder 
 ; 
 import 
  
 org.apache.beam.sdk.io.astra.db.mapping.AstraDbMapper 
 ; 
 import 
  
 org.apache.beam.sdk.io.astra.db.mapping.BeamRowDbMapperFactoryFn 
 ; 
 import 
  
 org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO 
 ; 
 import 
  
 org.apache.beam.sdk.io.gcp.bigquery.BigQueryUtils 
 ; 
 import 
  
 org.apache.beam.sdk.options.PipelineOptions 
 ; 
 import 
  
 org.apache.beam.sdk.options.PipelineOptionsFactory 
 ; 
 import 
  
 org.apache.beam.sdk.transforms.SerializableFunction 
 ; 
 import 
  
 org.apache.beam.sdk.values.Row 
 ; 
 import 
  
 org.apache.commons.lang3.StringUtils 
 ; 
 import 
  
 org.slf4j.Logger 
 ; 
 import 
  
 org.slf4j.LoggerFactory 
 ; 
 /** 
 * The {@link AstraDbToBigQuery} pipeline is a batch pipeline which ingests data from AstraDB and 
 * outputs the resulting records to BigQuery. 
 * 
 * <p>Check out <a 
 * href="https://github.com/GoogleCloudPlatform/DataflowTemplates/blob/main/v2/astradb-to-bigquery/README_AstraDB_to_BigQuery.md">README</a> 
 * for instructions on how to use or modify this template. 
 */ 
 @Template 
 ( 
  
 name 
  
 = 
  
 "AstraDB_To_BigQuery" 
 , 
  
 category 
  
 = 
  
 TemplateCategory 
 . 
 BATCH 
 , 
  
 displayName 
  
 = 
  
 "AstraDB to BigQuery" 
 , 
  
 description 
  
 = 
  
 { 
  
 "The AstraDB to BigQuery template is a batch pipeline that reads records from AstraDB and writes them to BigQuery." 
 , 
  
 "If the destination table doesn't exist in BigQuery, the pipeline creates a table with the following values:\n" 
  
 + 
  
 "- The `Dataset ID` is inherited from the Cassandra keyspace.\n" 
  
 + 
  
 "- The `Table ID` is inherited from the Cassandra table.\n" 
 , 
  
 "The schema of the destination table is inferred from the source Cassandra table.\n" 
  
 + 
  
 "- `List` and `Set` are mapped to BigQuery `REPEATED` fields.\n" 
  
 + 
  
 "- `Map` are mapped to BigQuery `RECORD` fields.\n" 
  
 + 
  
 "- All other types are mapped to BigQuery fields with the corresponding types.\n" 
  
 + 
  
 "- Cassandra user-defined types (UDTs) and tuple data types are not supported." 
  
 }, 
  
 optionsClass 
  
 = 
  
 AstraDbToBigQuery 
 . 
 Options 
 . 
 class 
 , 
  
 flexContainerName 
  
 = 
  
 "astradb-to-bigquery" 
 , 
  
 documentation 
  
 = 
  
 "https://cloud.google.com/dataflow/docs/guides/templates/provided/astradb-to-bigquery" 
 , 
  
 contactInformation 
  
 = 
  
 "https://cloud.google.com/support" 
 , 
  
 preview 
  
 = 
  
 true 
 , 
  
 requirements 
  
 = 
  
 { 
 "AstraDB account with a token" 
 }) 
 public 
  
 class 
 AstraDbToBigQuery 
  
 { 
  
 /** Logger for the class. */ 
  
 private 
  
 static 
  
 final 
  
 Logger 
  
 LOGGER 
  
 = 
  
 LoggerFactory 
 . 
 getLogger 
 ( 
 AstraDbToBigQuery 
 . 
 class 
 ); 
  
 /** If not provided, it is the default token range value. */ 
  
 public 
  
 static 
  
 final 
  
 int 
  
 DEFAULT_TOKEN_RANGE 
  
 = 
  
 18 
 ; 
  
 /** 
 * Options for the sample 
 * 
 * <p>Inherits standard configuration options. 
 */ 
  
 public 
  
 interface 
 Options 
  
 extends 
  
 PipelineOptions 
 , 
  
 AstraDbToBigQueryOptions 
 . 
 AstraDbSourceOptions 
 , 
  
 AstraDbToBigQueryOptions 
 . 
 BigQueryWriteOptions 
  
 {} 
  
 /** Main operations. */ 
  
 public 
  
 static 
  
 void 
  
 main 
 ( 
 String 
 [] 
  
 args 
 ) 
  
 { 
  
 UncaughtExceptionLogger 
 . 
 register 
 (); 
  
 LOGGER 
 . 
 info 
 ( 
 "Starting pipeline" 
 ); 
  
 try 
  
 { 
  
 Options 
  
 options 
  
 = 
  
 PipelineOptionsFactory 
 . 
 fromArgs 
 ( 
 args 
 ). 
 withValidation 
 (). 
 as 
 ( 
 Options 
 . 
 class 
 ); 
  
 LOGGER 
 . 
 debug 
 ( 
 "Pipeline Arguments (options) validated" 
 ); 
  
 // -------------------------------- 
  
 // AstraDbIO.Read<Row> 
  
 // -------------------------------- 
  
 // Credentials are read from secrets manager 
  
 AbstractMap 
 . 
 SimpleImmutableEntry<String 
 , 
  
 byte 
 [] 
>  
 astraCredentials 
  
 = 
  
 parseAstraCredentials 
 ( 
 options 
 ); 
  
 LOGGER 
 . 
 debug 
 ( 
 "Astra Credentials parsed" 
 ); 
  
 // Map Cassandra Table Schema into BigQuery Table Schema 
  
 SerializableFunction<AstraDbIO 
 . 
 Read 
< ? 
> , 
  
 TableSchema 
>  
 bigQuerySchemaFactory 
  
 = 
  
 new 
  
 AstraDbToBigQueryMappingFn 
 ( 
 options 
 . 
 getAstraKeyspace 
 (), 
  
 options 
 . 
 getAstraTable 
 ()); 
  
 LOGGER 
 . 
 debug 
 ( 
 "Schema Mapper has been initialized" 
 ); 
  
 // Map Cassandra Rows into (Apache) Beam Rows (DATA) 
  
 SerializableFunction<CqlSession 
 , 
  
 AstraDbMapper<Row> 
>  
 beamRowMapperFactory 
  
 = 
  
 new 
  
 BeamRowDbMapperFactoryFn 
 ( 
 options 
 . 
 getAstraKeyspace 
 (), 
  
 options 
 . 
 getAstraTable 
 ()); 
  
 LOGGER 
 . 
 debug 
 ( 
 "Row Mapper has been initialized" 
 ); 
  
 // Distribute reads across all available Cassandra nodes 
  
 int 
  
 minimalTokenRangesCount 
  
 = 
  
 ( 
 options 
 . 
 getMinTokenRangesCount 
 () 
  
 == 
  
 null 
 ) 
  
 ? 
  
 DEFAULT_TOKEN_RANGE 
  
 : 
  
 options 
 . 
 getMinTokenRangesCount 
 (); 
  
 // Source: AstraDb 
  
 AstraDbIO 
 . 
 Read<Row> 
  
 astraSource 
  
 = 
  
 AstraDbIO 
 . 
< Row>read 
 () 
  
 . 
 withToken 
 ( 
 astraCredentials 
 . 
 getKey 
 ()) 
  
 . 
 withSecureConnectBundle 
 ( 
 astraCredentials 
 . 
 getValue 
 ()) 
  
 . 
 withKeyspace 
 ( 
 options 
 . 
 getAstraKeyspace 
 ()) 
  
 . 
 withTable 
 ( 
 options 
 . 
 getAstraTable 
 ()) 
  
 . 
 withMinNumberOfSplits 
 ( 
 minimalTokenRangesCount 
 ) 
  
 . 
 withMapperFactoryFn 
 ( 
 beamRowMapperFactory 
 ) 
  
 . 
 withCoder 
 ( 
 SerializableCoder 
 . 
 of 
 ( 
 Row 
 . 
 class 
 )) 
  
 . 
 withEntity 
 ( 
 Row 
 . 
 class 
 ); 
  
 LOGGER 
 . 
 debug 
 ( 
 "AstraDb Source initialization [OK]" 
 ); 
  
 // -------------------------------- 
  
 //  BigQueryIO.Write<Row> 
  
 // -------------------------------- 
  
 TableReference 
  
 bqTableRef 
  
 = 
  
 parseBigQueryDestinationTable 
 ( 
 options 
 ); 
  
 createBigQueryDestinationTableIfNotExist 
 ( 
 options 
 , 
  
 bqTableRef 
 ); 
  
 LOGGER 
 . 
 debug 
 ( 
 "BigQuery Sink Table has been initialized" 
 ); 
  
 // Sink: BigQuery 
  
 BigQueryIO 
 . 
 Write<Row> 
  
 bigQuerySink 
  
 = 
  
 BigQueryIO 
 . 
< Row>write 
 () 
  
 . 
 to 
 ( 
 bqTableRef 
 ) 
  
 // Specialized function reading cassandra source table and mapping to BigQuery Schema 
  
 . 
 withSchema 
 ( 
 bigQuerySchemaFactory 
 . 
 apply 
 ( 
 astraSource 
 )) 
  
 // Provided by google, convert a Beam Row to a BigQuery TableRow 
  
 . 
 withFormatFunction 
 ( 
 row 
  
 - 
>  
 row 
  
 != 
  
 null 
  
 ? 
  
 BigQueryUtils 
 . 
 toTableRow 
 ( 
 row 
 ) 
  
 : 
  
 null 
 ) 
  
 // Table Will be created if not exist 
  
 . 
 withCreateDisposition 
 ( 
 BigQueryIO 
 . 
 Write 
 . 
 CreateDisposition 
 . 
 CREATE_IF_NEEDED 
 ) 
  
 . 
 withWriteDisposition 
 ( 
 BigQueryIO 
 . 
 Write 
 . 
 WriteDisposition 
 . 
 WRITE_APPEND 
 ); 
  
 LOGGER 
 . 
 debug 
 ( 
 "BigQuery Sink initialization [OK]" 
 ); 
  
 // -------------------------------- 
  
 //  Pipeline 
  
 // -------------------------------- 
  
 Pipeline 
  
 astraDbToBigQueryPipeline 
  
 = 
  
 Pipeline 
 . 
 create 
 ( 
 options 
 ); 
  
 astraDbToBigQueryPipeline 
  
 . 
 apply 
 ( 
 "Read From Astra" 
 , 
  
 astraSource 
 ) 
  
 . 
 apply 
 ( 
 "Write To BigQuery" 
 , 
  
 bigQuerySink 
 ); 
  
 astraDbToBigQueryPipeline 
 . 
 run 
 (); 
  
 } 
  
 finally 
  
 { 
  
 // Cassandra Connection is stateful and needs to be closed 
  
 CqlSessionHolder 
 . 
 cleanup 
 (); 
  
 } 
  
 } 
  
 /** 
 * Parse Astra Credentials from secrets in secret Manager. - SecretManagerUtils is not used as 
 * only applied to String secrets 
 * 
 * @param options pipeline options 
 * @return a pair with the token and the secure bundle 
 */ 
  
 private 
  
 static 
  
 AbstractMap 
 . 
 SimpleImmutableEntry<String 
 , 
  
 byte 
 [] 
>  
 parseAstraCredentials 
 ( 
  
 Options 
  
 options 
 ) 
  
 { 
  
 String 
  
 astraToken 
  
 = 
  
 options 
 . 
 getAstraToken 
 (); 
  
 if 
  
 ( 
 ! 
 astraToken 
 . 
 startsWith 
 ( 
 "AstraCS" 
 )) 
  
 { 
  
 astraToken 
  
 = 
  
 SecretManagerUtils 
 . 
 getSecret 
 ( 
 options 
 . 
 getAstraToken 
 ()); 
  
 } 
  
 LOGGER 
 . 
 info 
 ( 
 "Astra Token is parsed, value={}" 
 , 
  
 astraToken 
 . 
 substring 
 ( 
 0 
 , 
  
 10 
 ) 
  
 + 
  
 "..." 
 ); 
  
 /* 
 * Accessing the devops Api to retrieve the secure bundle. 
 */ 
  
 DatabaseClient 
  
 astraDbClient 
  
 = 
  
 new 
  
 DatabaseClient 
 ( 
 astraToken 
 , 
  
 options 
 . 
 getAstraDatabaseId 
 ()); 
  
 if 
  
 ( 
 ! 
 astraDbClient 
 . 
 exist 
 ()) 
  
 { 
  
 throw 
  
 new 
  
 RuntimeException 
 ( 
  
 "Astra Database does not exist, please check your Astra Token and Database ID" 
 ); 
  
 } 
  
 byte 
 [] 
  
 astraSecureBundle 
  
 = 
  
 astraDbClient 
 . 
 downloadDefaultSecureConnectBundle 
 (); 
  
 if 
  
 ( 
 ! 
 StringUtils 
 . 
 isEmpty 
 ( 
 options 
 . 
 getAstraDatabaseRegion 
 ())) 
  
 { 
  
 astraSecureBundle 
  
 = 
  
 astraDbClient 
 . 
 downloadSecureConnectBundle 
 ( 
 options 
 . 
 getAstraDatabaseRegion 
 ()); 
  
 } 
  
 LOGGER 
 . 
 info 
 ( 
 "Astra Bundle is parsed, length={}" 
 , 
  
 astraSecureBundle 
 . 
 length 
 ); 
  
 return 
  
 new 
  
 AbstractMap 
 . 
 SimpleImmutableEntry 
<> ( 
 astraToken 
 , 
  
 astraSecureBundle 
 ); 
  
 } 
  
 /** 
 * Create the Bog Query table Reference (provided or based on Cassandra table name). 
 * 
 * @param options pipeline options 
 * @return the big query table reference 
 */ 
  
 private 
  
 static 
  
 TableReference 
  
 parseBigQueryDestinationTable 
 ( 
 Options 
  
 options 
 ) 
  
 { 
  
 /* 
 * bigQueryOutputTableSpec argument is the Big Query table specification. This is parameter 
 * is optional. If not set, the table specification is built from the cassandra source table 
 * attributes: keyspace=dataset name, table=table name. 
 */ 
  
 String 
  
 bigQueryOutputTableSpec 
  
 = 
  
 options 
 . 
 getOutputTableSpec 
 (); 
  
 if 
  
 ( 
 StringUtils 
 . 
 isEmpty 
 ( 
 bigQueryOutputTableSpec 
 )) 
  
 { 
  
 bigQueryOutputTableSpec 
  
 = 
  
 options 
 . 
 getProject 
 () 
  
 + 
  
 ":" 
  
 + 
  
 options 
 . 
 getAstraKeyspace 
 () 
  
 + 
  
 "." 
  
 + 
  
 options 
 . 
 getAstraTable 
 (); 
  
 } 
  
 TableReference 
  
 bigQueryTableReference 
  
 = 
  
 BigQueryUtils 
 . 
 toTableReference 
 ( 
 bigQueryOutputTableSpec 
 ); 
  
 LOGGER 
 . 
 info 
 ( 
 "Big Query table spec has been set to {}" 
 , 
  
 bigQueryOutputTableSpec 
 ); 
  
 return 
  
 bigQueryTableReference 
 ; 
  
 } 
  
 /** 
 * Create destination dataset and tables if needed (schema mapped from Cassandra). 
 * 
 * @param options pipeline options 
 * @param bqTableRef big query table reference 
 */ 
  
 private 
  
 static 
  
 void 
  
 createBigQueryDestinationTableIfNotExist 
 ( 
  
 Options 
  
 options 
 , 
  
 TableReference 
  
 bqTableRef 
 ) 
  
 { 
  
 BigQuery 
  
 bigquery 
  
 = 
  
 BigQueryOptions 
 . 
 newBuilder 
 (). 
 setProjectId 
 ( 
 options 
 . 
 getProject 
 ()). 
 build 
 (). 
 getService 
 (); 
  
 if 
  
 ( 
 null 
  
 == 
  
 bigquery 
 . 
 getDataset 
 ( 
  
 DatasetId 
 . 
 of 
 ( 
 bqTableRef 
 . 
 getProjectId 
 (), 
  
 bqTableRef 
 . 
 getDatasetId 
 ()))) 
  
 { 
  
 LOGGER 
 . 
 info 
 ( 
  
 "Dataset was not found: creating DataSet {} in region {}" 
 , 
  
 bqTableRef 
 . 
 getDatasetId 
 (), 
  
 options 
 . 
 getWorkerRegion 
 ()); 
  
 bigquery 
 . 
 create 
 ( 
  
 DatasetInfo 
 . 
 newBuilder 
 ( 
 bqTableRef 
 . 
 getDatasetId 
 ()) 
  
 . 
 setLocation 
 ( 
 options 
 . 
 getWorkerRegion 
 ()) 
  
 . 
 build 
 ()); 
  
 LOGGER 
 . 
 debug 
 ( 
 "Dataset has been created [OK]" 
 ); 
  
 } 
  
 else 
  
 { 
  
 LOGGER 
 . 
 info 
 ( 
 "Dataset {} already exist" 
 , 
  
 bqTableRef 
 . 
 getDatasetId 
 ()); 
  
 } 
  
 } 
 }