Move autoCreateTable function to BigQueryWriter.

Author: Bingqin Zhou

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

@@ -142,8 +142,6 @@ private PartitionedTableId getRecordTable(SinkRecord record) {
 
     TableId baseTableId = topicsToBaseTableIds.get(record.topic());
 
-    maybeCreateTable(record, baseTableId);
-
     PartitionedTableId.Builder builder = new PartitionedTableId.Builder(baseTableId);
     if(usePartitionDecorator) {
 
@@ -161,20 +159,6 @@ private PartitionedTableId getRecordTable(SinkRecord record) {
     return builder.build();
   }
 
-  /**
-   * Create the table which doesn't exist in BigQuery for a (record's) topic when autoCreateTables config is set to true.
-   * @param record Kafka Sink Record to be streamed into BigQuery.
-   * @param baseTableId BaseTableId in BigQuery.
-   */
-  private void maybeCreateTable(SinkRecord record, TableId baseTableId) {
-    BigQuery bigQuery = getBigQuery();
-    boolean autoCreateTables = config.getBoolean(config.TABLE_CREATE_CONFIG);
-    if (autoCreateTables && bigQuery.getTable(baseTableId) == null) {
-      getSchemaManager(bigQuery).createTable(baseTableId, record.topic());
-      logger.info("Table {} does not exist, auto-created table for topic {}", baseTableId, record.topic());
-    }
-  }
-
   private RowToInsert getRecordRow(SinkRecord record) {
     Map<String, Object> convertedRecord = recordConverter.convertRecord(record, KafkaSchemaRecordType.VALUE);
     Optional<String> kafkaKeyFieldName = config.getKafkaKeyFieldName();
@@ -284,14 +268,17 @@ private SchemaManager getSchemaManager(BigQuery bigQuery) {
 
   private BigQueryWriter getBigQueryWriter() {
     boolean updateSchemas = config.getBoolean(config.SCHEMA_UPDATE_CONFIG);
+    boolean autoCreateTables = config.getBoolean(config.TABLE_CREATE_CONFIG);
     int retry = config.getInt(config.BIGQUERY_RETRY_CONFIG);
     long retryWait = config.getLong(config.BIGQUERY_RETRY_WAIT_CONFIG);
     BigQuery bigQuery = getBigQuery();
-    if (updateSchemas) {
+    if (updateSchemas || autoCreateTables) {
       return new AdaptiveBigQueryWriter(bigQuery,
                                         getSchemaManager(bigQuery),
                                         retry,
-                                        retryWait);
+                                        retryWait,
+                                        updateSchemas,
+                                        autoCreateTables);
     } else {
       return new SimpleBigQueryWriter(bigQuery, retry, retryWait);
     }

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

@@ -37,16 +37,18 @@
 import java.util.Map;
 
 /**
- * A {@link BigQueryWriter} capable of updating BigQuery table schemas.
+ * A {@link BigQueryWriter} capable of updating BigQuery table schemas and creating non-existed tables automatically.
  */
 public class AdaptiveBigQueryWriter extends BigQueryWriter {
   private static final Logger logger = LoggerFactory.getLogger(AdaptiveBigQueryWriter.class);
 
   // The maximum number of retries we will attempt to write rows after updating a BQ table schema.
-  private static final int AFTER_UPDATE_RETY_LIMIT = 5;
+  private static final int AFTER_UPDATE_RETRY_LIMIT = 5;
 
   private final BigQuery bigQuery;
   private final SchemaManager schemaManager;
+  private final boolean updateSchemas;
+  private final boolean autoCreateTables;
 
   /**
    * @param bigQuery Used to send write requests to BigQuery.
@@ -57,10 +59,14 @@ public class AdaptiveBigQueryWriter extends BigQueryWriter {
   public AdaptiveBigQueryWriter(BigQuery bigQuery,
                                 SchemaManager schemaManager,
                                 int retry,
-                                long retryWait) {
+                                long retryWait,
+                                boolean updateSchemas,
+                                boolean autoCreateTables) {
     super(retry, retryWait);
     this.bigQuery = bigQuery;
     this.schemaManager = schemaManager;
+    this.updateSchemas = updateSchemas;
+    this.autoCreateTables = autoCreateTables;
   }
 
   private boolean isTableMissingSchema(BigQueryException exception) {
@@ -69,6 +75,12 @@ private boolean isTableMissingSchema(BigQueryException exception) {
     return exception.getReason() != null && exception.getReason().equalsIgnoreCase("invalid");
   }
 
+  private boolean isTableNotExisted(BigQueryException exception) {
+    // If a table does not exist, it will raise a BigQueryException that the input is notFound
+    // Referring to Google Cloud Error Codes Doc: https://cloud.google.com/bigquery/docs/error-messages?hl=en
+    return exception.getReason() != null && exception.getReason().equalsIgnoreCase("notFound");
+  }
+
   /**
    * Sends the request to BigQuery, then checks the response to see if any errors have occurred. If
    * any have, and all errors can be blamed upon invalid columns in the rows sent, attempts to
@@ -89,11 +101,13 @@ public Map<Long, List<BigQueryError>> performWriteRequest(
       // Should only perform one schema update attempt; may have to continue insert attempts due to
       // BigQuery schema updates taking up to two minutes to take effect
       if (writeResponse.hasErrors()
-              && onlyContainsInvalidSchemaErrors(writeResponse.getInsertErrors())) {
+              && onlyContainsInvalidSchemaErrors(writeResponse.getInsertErrors()) && updateSchemas) {
         attemptSchemaUpdate(tableId, topic);
       }
     } catch (BigQueryException exception) {
-      if (isTableMissingSchema(exception)) {
+      if (isTableNotExisted(exception) && autoCreateTables) {
+        attemptTableCreate(tableId, topic);
+      } else if (isTableMissingSchema(exception) && updateSchemas) {
         attemptSchemaUpdate(tableId, topic);
       } else {
         throw exception;
@@ -117,10 +131,10 @@ && onlyContainsInvalidSchemaErrors(writeResponse.getInsertErrors())) {
         return writeResponse.getInsertErrors();
       }
       attemptCount++;
-      if (attemptCount >= AFTER_UPDATE_RETY_LIMIT) {
+      if (attemptCount >= AFTER_UPDATE_RETRY_LIMIT) {
         throw new BigQueryConnectException(
             "Failed to write rows after BQ schema update within "
-                + AFTER_UPDATE_RETY_LIMIT + " attempts for: " + tableId.getBaseTableId());
+                + AFTER_UPDATE_RETRY_LIMIT + " attempts for: " + tableId.getBaseTableId());
       }
     }
     logger.debug("table insertion completed successfully");
@@ -136,6 +150,16 @@ private void attemptSchemaUpdate(PartitionedTableId tableId, String topic) {
     }
   }
 
+  private void attemptTableCreate(PartitionedTableId tableId, String topic) {
+    try {
+      schemaManager.createTable(tableId.getBaseTableId(), topic);
+      logger.info("Table {} does not exist, auto-created table for topic {}", tableId.getBaseTableName(), topic);
+    } catch (BigQueryException exception) {
+      throw new BigQueryConnectException(
+              "Failed to create table " + tableId.getBaseTableName(), exception);
+    }
+  }
+
   /*
    * Currently, the only way to determine the cause of an insert all failure is by examining the map
    * object returned by the insertErrors() method of an insert all response. The only way to