Remove boolean retrieveKey from retrieveSchema.

Author: Bingqin Zhou

kcbq-api/src/main/java/com/wepay/kafka/connect/bigquery/api/SchemaRetriever.java

@@ -22,9 +22,10 @@ public interface SchemaRetriever {
    * Retrieve the most current schema for the given topic.
    * @param table The table that will be created.
    * @param topic The topic to retrieve a schema for.
+   * @param subjectType The type of kafka schema subject, either "value" or "key".
    * @return The Schema for the given table.
    */
-  public Schema retrieveSchema(TableId table, String topic, boolean retrieveKey);
+  public Schema retrieveSchema(TableId table, String topic, String subjectType);
 
   /**
    * Set the last seen schema for a given topic

kcbq-confluent/src/main/java/com/wepay/kafka/connect/bigquery/schemaregistry/schemaretriever/SchemaRegistrySchemaRetriever.java

@@ -61,8 +61,8 @@ public void configure(Map<String, String> properties) {
   }
 
   @Override
-  public Schema retrieveSchema(TableId table, String topic, boolean retrieveKey) {
-    String subject = getSubject(topic, retrieveKey);
+  public Schema retrieveSchema(TableId table, String topic, String subjectType) {
+    String subject = getSubject(topic, subjectType);
     try {
       logger.debug("Retrieving schema information for topic {} with subject {}", topic, subject);
       SchemaMetadata latestSchemaMetadata = schemaRegistryClient.getLatestSchemaMetadata(subject);
@@ -80,10 +80,7 @@ public Schema retrieveSchema(TableId table, String topic, boolean retrieveKey) {
   @Override
   public void setLastSeenSchema(TableId table, String topic, Schema schema) { }
 
-  private String getSubject(String topic, boolean retrieveKey) {
-    if (retrieveKey) {
-      return topic + "-key";
-    }
-    return topic + "-value";
+  private String getSubject(String topic, String subjectType) {
+    return topic + "-" + subjectType;
   }
 }

kcbq-confluent/src/test/java/com/wepay/kafka/connect/bigquery/schemaregistry/schemaretriever/SchemaRegistrySchemaRetrieverTest.java

@@ -42,7 +42,7 @@ public void testRetrieveSchema() throws Exception {
     Schema expectedKafkaConnectSchema =
         SchemaBuilder.struct().field("f1", Schema.STRING_SCHEMA).name("testrecord").build();
 
-    assertEquals(expectedKafkaConnectSchema, testSchemaRetriever.retrieveSchema(table, testTopic, false));
-    assertEquals(expectedKafkaConnectSchema, testSchemaRetriever.retrieveSchema(table, testTopic, true));
+    assertEquals(expectedKafkaConnectSchema, testSchemaRetriever.retrieveSchema(table, testTopic, "value"));
+    assertEquals(expectedKafkaConnectSchema, testSchemaRetriever.retrieveSchema(table, testTopic, "key"));
   }
 }

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

@@ -18,6 +18,8 @@
  */
 public class SchemaManager {
   private static final Logger logger = LoggerFactory.getLogger(SchemaManager.class);
+  private static final String VALUE = "value";
+  private static final String KEY = "key";
 
   private final SchemaRetriever schemaRetriever;
   private final SchemaConverter<com.google.cloud.bigquery.Schema> schemaConverter;
@@ -52,8 +54,8 @@ public SchemaManager(
    * @param topic The Kafka topic used to determine the schema.
    */
   public void createTable(TableId table, String topic) {
-    Schema kafkaValueSchema = schemaRetriever.retrieveSchema(table, topic, false);
-    Schema kafkaKeySchema = includeKafkaKey ? schemaRetriever.retrieveSchema(table, topic, true) : null;
+    Schema kafkaValueSchema = schemaRetriever.retrieveSchema(table, topic, VALUE);
+    Schema kafkaKeySchema = includeKafkaKey ? schemaRetriever.retrieveSchema(table, topic, KEY) : null;
     bigQuery.create(constructTableInfo(table, kafkaKeySchema, kafkaValueSchema));
   }
 
@@ -63,8 +65,8 @@ public void createTable(TableId table, String topic) {
    * @param topic The Kafka topic used to determine the schema.
    */
   public void updateSchema(TableId table, String topic) {
-    Schema kafkaValueSchema = schemaRetriever.retrieveSchema(table, topic, false);
-    Schema kafkaKeySchema = includeKafkaKey ? schemaRetriever.retrieveSchema(table, topic, true) : null;
+    Schema kafkaValueSchema = schemaRetriever.retrieveSchema(table, topic, VALUE);
+    Schema kafkaKeySchema = includeKafkaKey ? schemaRetriever.retrieveSchema(table, topic, KEY) : null;
     TableInfo tableInfo = constructTableInfo(table, kafkaKeySchema, kafkaValueSchema);
     logger.info("Attempting to update table `{}` with schema {}",
         table, tableInfo.getDefinition().getSchema());

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

@@ -8,6 +8,9 @@
 import java.util.HashMap;
 import java.util.Map;
 
+/**
+ * Helper class to construct kafka data schema and kafka data record.
+ */
 public class KafkaDataConverter {
 
     public static final String KAFKA_DATA_FIELD_NAME = "kafkaData";

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

@@ -40,7 +40,7 @@ public void configure(Map<String, String> properties) {
   }
 
   @Override
-  public Schema retrieveSchema(TableId table, String topic, boolean retrieveKey) {
+  public Schema retrieveSchema(TableId table, String topic, String subjectType) {
     String tableName = table.getTable();
     Schema schema = schemaCache.get(getCacheKey(tableName, topic));
     if (schema != null) {

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

@@ -61,7 +61,7 @@ public void configure(Map<String, String> properties) {
     }
 
     @Override
-    public Schema retrieveSchema(TableId table, String topic, boolean retrieveKey) {
+    public Schema retrieveSchema(TableId table, String topic, String subjectType) {
       // Shouldn't be called
       return null;
     }

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

@@ -25,7 +25,7 @@ public void testRetrieveSchemaWhenNoLastSeenSchemaReturnsEmptyStructSchema() {
     final TableId tableId = getTableId("testTable", "testDataset");
     SchemaRetriever retriever = new MemorySchemaRetriever();
     retriever.configure(new HashMap<>());
-    Assert.assertEquals(retriever.retrieveSchema(tableId, topic, false), SchemaBuilder.struct().build());
+    Assert.assertEquals(retriever.retrieveSchema(tableId, topic, "value"), SchemaBuilder.struct().build());
   }
 
   @Test
@@ -38,7 +38,7 @@ public void testRetrieveSchemaWhenLastSeenExistsSucceeds() {
     Schema expectedSchema = Schema.OPTIONAL_FLOAT32_SCHEMA;
     retriever.setLastSeenSchema(tableId, topic, expectedSchema);
 
-    Assert.assertEquals(retriever.retrieveSchema(tableId, topic, true), expectedSchema);
+    Assert.assertEquals(retriever.retrieveSchema(tableId, topic, "key"), expectedSchema);
   }
 
   @Test
@@ -56,8 +56,8 @@ public void testRetrieveSchemaWithMultipleSchemasSucceeds() {
     retriever.setLastSeenSchema(intTableId, intSchemaTopic, expectedIntSchema);
 
     Assert.assertEquals(
-        retriever.retrieveSchema(floatTableId, floatSchemaTopic, false), expectedFloatSchema);
-    Assert.assertEquals(retriever.retrieveSchema(intTableId, intSchemaTopic, true), expectedIntSchema);
+        retriever.retrieveSchema(floatTableId, floatSchemaTopic, "value"), expectedFloatSchema);
+    Assert.assertEquals(retriever.retrieveSchema(intTableId, intSchemaTopic, "key"), expectedIntSchema);
   }
 
   @Test
@@ -72,6 +72,6 @@ public void testRetrieveSchemaRetrievesLastSeenSchema() {
     retriever.setLastSeenSchema(tableId, intSchemaTopic, firstSchema);
     retriever.setLastSeenSchema(tableId, intSchemaTopic, secondSchema);
 
-    Assert.assertEquals(retriever.retrieveSchema(tableId, intSchemaTopic, false), secondSchema);
+    Assert.assertEquals(retriever.retrieveSchema(tableId, intSchemaTopic, "value"), secondSchema);
   }
 }