Auto-create tables for GCS Load mode.

Author: Bingqin Zhou

kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/BigQuerySinkTask.java

@@ -201,7 +201,8 @@ public void put(Collection<SinkRecord> records) {
         if (!tableWriterBuilders.containsKey(table)) {
           TableWriterBuilder tableWriterBuilder;
           if (config.getList(config.ENABLE_BATCH_CONFIG).contains(record.topic())) {
-            String gcsBlobName = record.topic() + "_" + uuid + "_" + Instant.now().toEpochMilli();
+            String topic = record.topic();
+            String gcsBlobName = topic + "_" + uuid + "_" + Instant.now().toEpochMilli();
             String gcsFolderName = config.getString(config.GCS_FOLDER_NAME_CONFIG);
             if (gcsFolderName != null && !"".equals(gcsFolderName)) {
               gcsBlobName = gcsFolderName + "/" + gcsBlobName;
@@ -211,6 +212,7 @@ public void put(Collection<SinkRecord> records) {
                 table.getBaseTableId(),
                 config.getString(config.GCS_BUCKET_NAME_CONFIG),
                 gcsBlobName,
+                topic,
                 recordConverter);
           } else {
             tableWriterBuilder =
@@ -299,10 +301,13 @@ private GCSToBQWriter getGcsWriter() {
     BigQuery bigQuery = getBigQuery();
     int retry = config.getInt(config.BIGQUERY_RETRY_CONFIG);
     long retryWait = config.getLong(config.BIGQUERY_RETRY_WAIT_CONFIG);
+    boolean autoCreateTables = config.getBoolean(config.TABLE_CREATE_CONFIG);
     return new GCSToBQWriter(getGcs(),
                          bigQuery,
+                         getSchemaManager(bigQuery),
                          retry,
-                         retryWait);
+                         retryWait,
+                         autoCreateTables);
   }
 
   @Override

kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/write/batch/GCSBatchTableWriter.java

@@ -45,6 +45,7 @@ public class GCSBatchTableWriter implements Runnable {
 
   private final String bucketName;
   private final String blobName;
+  private final String topic;
 
   private final List<RowToInsert> rows;
   private final GCSToBQWriter writer;
@@ -61,10 +62,12 @@ private GCSBatchTableWriter(List<RowToInsert> rows,
                               GCSToBQWriter writer,
                               TableId tableId,
                               String bucketName,
-                              String baseBlobName) {
+                              String baseBlobName,
+                              String topic) {
     this.tableId = tableId;
     this.bucketName = bucketName;
     this.blobName = baseBlobName;
+    this.topic = topic;
 
     this.rows = rows;
     this.writer = writer;
@@ -73,7 +76,7 @@ private GCSBatchTableWriter(List<RowToInsert> rows,
   @Override
   public void run() {
     try {
-      writer.writeRows(rows, tableId, bucketName, blobName);
+      writer.writeRows(rows, tableId, bucketName, blobName, topic);
     } catch (ConnectException ex) {
       throw new ConnectException("Failed to write rows to GCS", ex);
     } catch (InterruptedException ex) {
@@ -87,6 +90,7 @@ public void run() {
   public static class Builder implements TableWriterBuilder {
     private final String bucketName;
     private String blobName;
+    private String topic;
 
     private final TableId tableId;
 
@@ -107,10 +111,12 @@ public Builder(GCSToBQWriter writer,
                    TableId tableId,
                    String gcsBucketName,
                    String gcsBlobName,
+                   String topic,
                    RecordConverter<Map<String, Object>> recordConverter) {
 
       this.bucketName = gcsBucketName;
       this.blobName = gcsBlobName;
+      this.topic = topic;
 
       this.tableId = tableId;
 
@@ -133,7 +139,7 @@ public void addRow(RowToInsert rowToInsert) {
     }
 
     public GCSBatchTableWriter build() {
-      return new GCSBatchTableWriter(rows, writer, tableId, bucketName, blobName);
+      return new GCSBatchTableWriter(rows, writer, tableId, bucketName, blobName, topic);
     }
   }
 }

kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/write/row/AdaptiveBigQueryWriter.java

@@ -19,6 +19,7 @@
 
 
 import com.google.cloud.bigquery.BigQuery;
+import com.google.cloud.bigquery.TableId;
 import com.google.cloud.bigquery.BigQueryError;
 import com.google.cloud.bigquery.BigQueryException;
 import com.google.cloud.bigquery.InsertAllRequest;
@@ -109,7 +110,7 @@ && onlyContainsInvalidSchemaErrors(writeResponse.getInsertErrors()) && updateSch
     } catch (BigQueryException exception) {
       // Should only perform one table creation attempt.
       if (isTableNotExisted(exception) && autoCreateTables && bigQuery.getTable(tableId.getBaseTableId()) == null) {
-        attemptTableCreate(tableId, topic);
+        attemptTableCreate(tableId.getBaseTableId(), topic);
       } else if (isTableMissingSchema(exception) && updateSchemas) {
         attemptSchemaUpdate(tableId, topic);
       } else {
@@ -158,13 +159,13 @@ private void attemptSchemaUpdate(PartitionedTableId tableId, String topic) {
     }
   }
 
-  private void attemptTableCreate(PartitionedTableId tableId, String topic) {
+  private void attemptTableCreate(TableId tableId, String topic) {
     try {
-      schemaManager.createTable(tableId.getBaseTableId(), topic);
-      logger.info("Table {} does not exist, auto-created table for topic {}", tableId.getBaseTableName(), topic);
+      schemaManager.createTable(tableId, topic);
+      logger.info("Table {} does not exist, auto-created table for topic {}", tableId, topic);
     } catch (BigQueryException exception) {
       throw new BigQueryConnectException(
-              "Failed to create table " + tableId.getBaseTableName(), exception);
+              "Failed to create table " + tableId, exception);
     }
   }
 

kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/write/row/GCSToBQWriter.java

@@ -19,6 +19,7 @@
 
 
 import com.google.cloud.bigquery.BigQuery;
+import com.google.cloud.bigquery.BigQueryException;
 import com.google.cloud.bigquery.InsertAllRequest.RowToInsert;
 import com.google.cloud.bigquery.TableId;
 import com.google.cloud.storage.Blob;
@@ -28,6 +29,8 @@
 import com.google.cloud.storage.StorageException;
 import com.google.gson.Gson;
 
+import com.wepay.kafka.connect.bigquery.SchemaManager;
+import com.wepay.kafka.connect.bigquery.exception.BigQueryConnectException;
 import com.wepay.kafka.connect.bigquery.exception.GCSConnectException;
 
 import org.apache.kafka.connect.errors.ConnectException;
@@ -53,12 +56,16 @@ public class GCSToBQWriter {
 
   private final BigQuery bigQuery;
 
+  private final SchemaManager schemaManager;
+
   private static final int WAIT_MAX_JITTER = 1000;
 
   private static final Random random = new Random();
 
   private int retries;
   private long retryWaitMs;
+  private boolean autoCreateTables;
+
 
   public static final String GCS_METADATA_TABLE_KEY = "sinkTable";
 
@@ -71,13 +78,17 @@ public class GCSToBQWriter {
    */
   public GCSToBQWriter(Storage storage,
                        BigQuery bigQuery,
+                       SchemaManager schemaManager,
                        int retries,
-                       long retryWaitMs) {
+                       long retryWaitMs,
+                       boolean autoCreateTables) {
     this.storage = storage;
     this.bigQuery = bigQuery;
+    this.schemaManager = schemaManager;
 
     this.retries = retries;
     this.retryWaitMs = retryWaitMs;
+    this.autoCreateTables = autoCreateTables;
   }
 
   /**
@@ -92,7 +103,8 @@ public GCSToBQWriter(Storage storage,
   public void writeRows(List<RowToInsert> rows,
                         TableId tableId,
                         String bucketName,
-                        String blobName) throws InterruptedException {
+                        String blobName,
+                        String topic) throws InterruptedException {
 
     // Get Source URI
     BlobId blobId = BlobId.of(bucketName, blobName);
@@ -103,9 +115,8 @@ public void writeRows(List<RowToInsert> rows,
 
     // Check if the table specified exists
     // This error shouldn't be thrown. All tables should be created by the connector at startup
-    if (bigQuery.getTable(tableId) == null) {
-      throw new ConnectException(
-          String.format("Table with TableId %s does not exist.", tableId.getTable()));
+    if (autoCreateTables && bigQuery.getTable(tableId) == null) {
+      attemptTableCreate(tableId, topic);
     }
 
     int attemptCount = 0;
@@ -184,4 +195,14 @@ private String toJson(List<RowToInsert> rows) {
   private void waitRandomTime() throws InterruptedException {
     Thread.sleep(retryWaitMs + random.nextInt(WAIT_MAX_JITTER));
   }
+
+  private void attemptTableCreate(TableId tableId, String topic) {
+    try {
+      schemaManager.createTable(tableId, topic);
+      logger.info("Table {} does not exist, auto-created table for topic {}", tableId, topic);
+    } catch (BigQueryException exception) {
+      throw new BigQueryConnectException(
+              "Failed to create table " + tableId, exception);
+    }
+  }
 }