Make Kafka data field and key name configurable.

Author: Bingqin Zhou

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

@@ -104,7 +104,10 @@ private SchemaManager getSchemaManager(BigQuery bigQuery) {
     SchemaConverter<com.google.cloud.bigquery.Schema> schemaConverter = config.getSchemaConverter();
     boolean includeKafkaKey = config.getBoolean(config.INCLUDE_KAFKA_KEY_CONFIG);
     boolean includeKafkaData = config.getBoolean(config.INCLUDE_KAFKA_DATA_CONFIG);
-    return new SchemaManager(schemaRetriever, schemaConverter, bigQuery, includeKafkaKey, includeKafkaData);
+    String kafkaKeyFieldName = config.getString(config.KAFKA_KEY_FIELD_NAME_CONFIG);
+    String kafkaDataFieldName = config.getString(config.KAFKA_DATA_FIELD_NAME_CONFIG);
+    return new SchemaManager(schemaRetriever, schemaConverter, bigQuery,
+            includeKafkaKey, includeKafkaData, kafkaKeyFieldName, kafkaDataFieldName);
   }
 
   private void ensureExistingTables(

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

@@ -28,7 +28,7 @@
 import com.wepay.kafka.connect.bigquery.api.SchemaRetriever;
 import com.wepay.kafka.connect.bigquery.config.BigQuerySinkConfig;
 import com.wepay.kafka.connect.bigquery.config.BigQuerySinkTaskConfig;
-import com.wepay.kafka.connect.bigquery.convert.KafkaDataConverter;
+import com.wepay.kafka.connect.bigquery.convert.KafkaDataBuilder;
 import com.wepay.kafka.connect.bigquery.convert.RecordConverter;
 import com.wepay.kafka.connect.bigquery.convert.SchemaConverter;
 import com.wepay.kafka.connect.bigquery.exception.SinkConfigConnectException;
@@ -153,10 +153,10 @@ private PartitionedTableId getRecordTable(SinkRecord record) {
   private RowToInsert getRecordRow(SinkRecord record) {
     Map<String, Object> convertedRecord = recordConverter.convertRecord(record, false);
     if (config.getBoolean(config.INCLUDE_KAFKA_KEY_CONFIG)) {
-      convertedRecord.put(SchemaManager.KAFKA_KEY_FIELD_NAME, recordConverter.convertRecord(record, true));
+      convertedRecord.put(config.KAFKA_KEY_FIELD_NAME_CONFIG, recordConverter.convertRecord(record, true));
     }
     if (config.getBoolean(config.INCLUDE_KAFKA_DATA_CONFIG)) {
-      convertedRecord.put(KafkaDataConverter.KAFKA_DATA_FIELD_NAME, KafkaDataConverter.getKafkaDataRecord(record));
+      convertedRecord.put(config.KAFKA_DATA_FIELD_NAME_CONFIG, KafkaDataBuilder.getKafkaDataRecord(record));
     }
     if (config.getBoolean(config.SANITIZE_FIELD_NAME_CONFIG)) {
       convertedRecord = FieldNameSanitizer.replaceInvalidKeys(convertedRecord);
@@ -249,7 +249,10 @@ private SchemaManager getSchemaManager(BigQuery bigQuery) {
         config.getSchemaConverter();
     boolean includeKafkaKey = config.getBoolean(config.INCLUDE_KAFKA_KEY_CONFIG);
     boolean includeKafkaData = config.getBoolean(config.INCLUDE_KAFKA_DATA_CONFIG);
-    return new SchemaManager(schemaRetriever, schemaConverter, bigQuery, includeKafkaKey, includeKafkaData);
+    String kafkaKeyFieldName = config.getString(config.KAFKA_KEY_FIELD_NAME_CONFIG);
+    String kafkaDataFieldName = config.getString(config.KAFKA_DATA_FIELD_NAME_CONFIG);
+    return new SchemaManager(schemaRetriever, schemaConverter, bigQuery,
+            includeKafkaKey, includeKafkaData, kafkaKeyFieldName, kafkaDataFieldName);
   }
 
   private BigQueryWriter getBigQueryWriter() {

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

@@ -3,7 +3,7 @@
 
 import com.google.cloud.bigquery.*;
 import com.wepay.kafka.connect.bigquery.api.SchemaRetriever;
-import com.wepay.kafka.connect.bigquery.convert.KafkaDataConverter;
+import com.wepay.kafka.connect.bigquery.convert.KafkaDataBuilder;
 import com.wepay.kafka.connect.bigquery.convert.SchemaConverter;
 
 import org.apache.kafka.connect.data.Schema;
@@ -26,8 +26,8 @@ public class SchemaManager {
   private final BigQuery bigQuery;
   private final boolean includeKafkaKey;
   private final boolean includeKafkaData;
-
-  public static final String KAFKA_KEY_FIELD_NAME = "kafkaKey";
+  private final String kafkaKeyFieldName;
+  private final String kafkaDataFieldName;
 
   /**
    * @param schemaRetriever Used to determine the Kafka Connect Schema that should be used for a
@@ -40,12 +40,16 @@ public SchemaManager(
       SchemaConverter<com.google.cloud.bigquery.Schema> schemaConverter,
       BigQuery bigQuery,
       boolean includeKafkaKey,
-      boolean includeKafkaData) {
+      boolean includeKafkaData,
+      String kafkaKeyFieldName,
+      String kafkaDataFieldName) {
     this.schemaRetriever = schemaRetriever;
     this.schemaConverter = schemaConverter;
     this.bigQuery = bigQuery;
     this.includeKafkaKey = includeKafkaKey;
     this.includeKafkaData = includeKafkaData;
+    this.kafkaKeyFieldName = kafkaKeyFieldName;
+    this.kafkaDataFieldName = kafkaDataFieldName;
   }
 
   /**
@@ -94,12 +98,12 @@ private com.google.cloud.bigquery.Schema getBigQuerySchema(Schema kafkaKeySchema
       allFields.addAll(valueSchema.getFields());
       if (includeKafkaKey) {
           com.google.cloud.bigquery.Schema keySchema = schemaConverter.convertSchema(kafkaKeySchema);
-          Field kafkaKeyField = Field.newBuilder(KAFKA_KEY_FIELD_NAME, LegacySQLTypeName.RECORD, keySchema.getFields())
+          Field kafkaKeyField = Field.newBuilder(kafkaKeyFieldName, LegacySQLTypeName.RECORD, keySchema.getFields())
                   .setMode(Field.Mode.NULLABLE).build();
           allFields.add(kafkaKeyField);
       }
       if (includeKafkaData) {
-          Field kafkaDataField = KafkaDataConverter.getKafkaDataField();
+          Field kafkaDataField = KafkaDataBuilder.getKafkaDataField(kafkaDataFieldName);
           allFields.add(kafkaDataField);
       }
       return com.google.cloud.bigquery.Schema.of(allFields);

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

@@ -183,6 +183,18 @@ public class BigQuerySinkConfig extends AbstractConfig {
   public static final String INCLUDE_KAFKA_KEY_DOC =
           "Whether to include an extra block containing fields in Kafka key.";
 
+  public static final String KAFKA_KEY_FIELD_NAME_CONFIG =        "kafkaKeyFieldName";
+  private static final ConfigDef.Type KAFKA_KEY_FIELD_NAME_TYPE = ConfigDef.Type.STRING;
+  public static final String KAFKA_KEY_FIELD_NAME_DEFAULT =       "kafkaKey";
+  private static final ConfigDef.Importance KAFKA_KEY_FIELD_NAME_IMPORTANCE = ConfigDef.Importance.LOW;
+  private static final String KAFKA_KEY_FIELD_NAME_DOC = "The name of the field of Kafka key.";
+
+  public static final String KAFKA_DATA_FIELD_NAME_CONFIG =        "kafkaDataFieldName";
+  private static final ConfigDef.Type KAFKA_DATA_FIELD_NAME_TYPE = ConfigDef.Type.STRING;
+  public static final String KAFKA_DATA_FIELD_NAME_DEFAULT =       "kafkaData";
+  private static final ConfigDef.Importance KAFKA_DATA_FIELD_NAME_IMPORTANCE = ConfigDef.Importance.LOW;
+  private static final String KAFKA_DATA_FIELD_NAME_DOC = "The name of the field of Kafka Data.";
+
   public static final String AVRO_DATA_CACHE_SIZE_CONFIG =                 "avroDataCacheSize";
   private static final ConfigDef.Type AVRO_DATA_CACHE_SIZE_TYPE =          ConfigDef.Type.INT;
   public static final Integer AVRO_DATA_CACHE_SIZE_DEFAULT =               100;
@@ -307,6 +319,18 @@ public class BigQuerySinkConfig extends AbstractConfig {
             INCLUDE_KAFKA_KEY_DEFAULT,
             INCLUDE_KAFKA_KEY_IMPORTANCE,
             INCLUDE_KAFKA_KEY_DOC
+        ).define(
+            KAFKA_KEY_FIELD_NAME_CONFIG,
+            KAFKA_KEY_FIELD_NAME_TYPE,
+            KAFKA_KEY_FIELD_NAME_DEFAULT,
+            KAFKA_KEY_FIELD_NAME_IMPORTANCE,
+            KAFKA_KEY_FIELD_NAME_DOC
+        ).define(
+            KAFKA_DATA_FIELD_NAME_CONFIG,
+            KAFKA_DATA_FIELD_NAME_TYPE,
+            KAFKA_DATA_FIELD_NAME_DEFAULT,
+            KAFKA_DATA_FIELD_NAME_IMPORTANCE,
+            KAFKA_DATA_FIELD_NAME_DOC
         ).define(
             AVRO_DATA_CACHE_SIZE_CONFIG,
             AVRO_DATA_CACHE_SIZE_TYPE,

kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/convert/KafkaDataConverter.java renamed to kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/convert/KafkaDataBuilder.java

@@ -11,15 +11,14 @@
 /**
  * Helper class to construct kafka data schema and kafka data record.
  */
-public class KafkaDataConverter {
+public class KafkaDataBuilder {
 
-    public static final String KAFKA_DATA_FIELD_NAME = "kafkaData";
     public static final String KAFKA_DATA_TOPIC_FIELD_NAME = "topic";
     public static final String KAFKA_DATA_PARTITION_FIELD_NAME = "partition";
     public static final String KAFKA_DATA_OFFSET_FIELD_NAME = "offset";
     public static final String KAFKA_DATA_INSERT_TIME_FIELD_NAME = "insertTime";
 
-    public static Field getKafkaDataField() {
+    public static Field getKafkaDataField(String kafkaDataFieldName) {
         Field topicField = com.google.cloud.bigquery.Field.of(KAFKA_DATA_TOPIC_FIELD_NAME, LegacySQLTypeName.STRING);
         Field partitionField = com.google.cloud.bigquery.Field.of(KAFKA_DATA_PARTITION_FIELD_NAME, LegacySQLTypeName.INTEGER);
         Field offsetField = com.google.cloud.bigquery.Field.of(KAFKA_DATA_OFFSET_FIELD_NAME, LegacySQLTypeName.INTEGER);
@@ -28,7 +27,7 @@ public static Field getKafkaDataField() {
                 .setMode(com.google.cloud.bigquery.Field.Mode.NULLABLE);
         Field insertTimeField = insertTimeBuilder.build();
 
-        return Field.newBuilder(KAFKA_DATA_FIELD_NAME, LegacySQLTypeName.RECORD,
+        return Field.newBuilder(kafkaDataFieldName, LegacySQLTypeName.RECORD,
                 topicField, partitionField, offsetField, insertTimeField)
                 .setMode(com.google.cloud.bigquery.Field.Mode.NULLABLE).build();
     }

kcbq-connector/src/test/java/com/wepay/kafka/connect/bigquery/SchemaManagerTest.java

@@ -54,7 +54,9 @@ public void testBQTableDescription() {
                                                     mockSchemaConverter,
                                                     mockBigQuery,
                                                     true,
-                                                    true);
+                                                    true,
+                                                    "kafkaKey",
+                                                     "kafkaData");
 
     Schema mockKafkaSchema = mock(Schema.class);
     // we would prefer to mock this class, but it is final.

kcbq-connector/src/test/java/com/wepay/kafka/connect/bigquery/convert/KafkaDataConverterTest.java

@@ -14,7 +14,7 @@
 
 public class KafkaDataConverterTest {
 
-    public static final String KAFKA_DATA_FIELD_NAME = "kafkaData";
+    public static final String kafkaDataFieldName = "kafkaData";
     private static final String kafkaDataTopicName = "topic";
     private static final String kafkaDataPartitionName = "partition";
     private static final String kafkaDataOffsetName = "offset";
@@ -33,7 +33,7 @@ public void testGetKafkaDataRecord() {
         expectedKafkaDataFields.put(kafkaDataOffsetName, kafkaDataOffsetValue);
 
         SinkRecord record = new SinkRecord(kafkaDataTopicValue, kafkaDataPartitionValue, null, null, null, null, kafkaDataOffsetValue);
-        Map<String, Object> actualKafkaDataFields = KafkaDataConverter.getKafkaDataRecord(record);
+        Map<String, Object> actualKafkaDataFields = KafkaDataBuilder.getKafkaDataRecord(record);
 
         assertTrue(actualKafkaDataFields.containsKey(kafkaDataInsertTimeName));
         assertTrue(actualKafkaDataFields.get(kafkaDataInsertTimeName) instanceof Double);
@@ -52,15 +52,15 @@ public void testGetKafkaDataField() {
                 .setMode(Field.Mode.NULLABLE)
                 .build();
 
-        Field expectedBigQuerySchema = Field.newBuilder(KAFKA_DATA_FIELD_NAME,
+        Field expectedBigQuerySchema = Field.newBuilder(kafkaDataFieldName,
                 LegacySQLTypeName.RECORD,
                 topicField,
                 partitionField,
                 offsetField,
                 insertTimeField)
                 .setMode(Field.Mode.NULLABLE)
                 .build();
-        Field actualBigQuerySchema = KafkaDataConverter.getKafkaDataField();
+        Field actualBigQuerySchema = KafkaDataBuilder.getKafkaDataField(kafkaDataFieldName);
         assertEquals(expectedBigQuerySchema, actualBigQuerySchema);
     }
 }