fix: populate populateLogicalConverterRegistry in Task

Claudenw · web-flow · commit 78f5a4d0a211 · 2025-12-05T11:13:18.000+01:00
diff --git a/kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/BigQuerySinkConnector.java b/kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/BigQuerySinkConnector.java
@@ -73,8 +73,6 @@ public void start(Map<String, String> properties) {
     logger.trace("connector.start()");
     configProperties = properties;
     config = new BigQuerySinkConfig(properties);
-    DebeziumLogicalConverters.initialize(config);
-    KafkaLogicalConverters.initialize(config);
   }
 
   @Override
diff --git a/kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/BigQuerySinkTask.java b/kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/BigQuerySinkTask.java
@@ -37,6 +37,8 @@
 import com.wepay.kafka.connect.bigquery.config.BigQuerySinkConfig;
 import com.wepay.kafka.connect.bigquery.config.BigQuerySinkTaskConfig;
 import com.wepay.kafka.connect.bigquery.convert.SchemaConverter;
+import com.wepay.kafka.connect.bigquery.convert.logicaltype.DebeziumLogicalConverters;
+import com.wepay.kafka.connect.bigquery.convert.logicaltype.KafkaLogicalConverters;
 import com.wepay.kafka.connect.bigquery.exception.BigQueryStorageWriteApiConnectException;
 import com.wepay.kafka.connect.bigquery.exception.ConversionConnectException;
 import com.wepay.kafka.connect.bigquery.utils.PartitionedTableId;
@@ -419,11 +421,17 @@ boolean isRunning() {
     return !stopped;
   }
 
+  private static void populateLogicalConverterRegistry(BigQuerySinkTaskConfig config) {
+    DebeziumLogicalConverters.initialize(config);
+    KafkaLogicalConverters.initialize(config);
+  }
+
   @Override
   public void start(Map<String, String> properties) {
     logger.trace("task.start()");
     stopped = false;
     config = new BigQuerySinkTaskConfig(properties);
+    populateLogicalConverterRegistry(config);
     autoCreateTables = config.getBoolean(BigQuerySinkConfig.TABLE_CREATE_CONFIG);
 
     useStorageApi = config.getBoolean(BigQuerySinkConfig.USE_STORAGE_WRITE_API_CONFIG);
diff --git a/kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/convert/logicaltype/DebeziumLogicalConverters.java b/kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/convert/logicaltype/DebeziumLogicalConverters.java
@@ -25,6 +25,7 @@
 
 import com.google.cloud.bigquery.Field;
 import com.google.cloud.bigquery.LegacySQLTypeName;
+import com.google.common.annotations.VisibleForTesting;
 import com.wepay.kafka.connect.bigquery.config.BigQuerySinkConfig;
 import com.wepay.kafka.connect.bigquery.exception.ConversionConnectException;
 import io.debezium.data.VariableScaleDecimal;
@@ -52,14 +53,32 @@ public class DebeziumLogicalConverters {
   private static final int MICROS_IN_SEC = 1000000;
   private static final int MICROS_IN_MILLI = 1000;
 
+  /**
+   * Initialize the LogicalConverterRegistry with the DebeziumLogicalConverters.
+   *
+   * @param config the configuration to use.
+   */
   public static void initialize(final BigQuerySinkConfig config) {
-    LogicalConverterRegistry.register(Date.SCHEMA_NAME, new DateConverter());
-    LogicalConverterRegistry.register(MicroTime.SCHEMA_NAME, new MicroTimeConverter());
-    LogicalConverterRegistry.register(MicroTimestamp.SCHEMA_NAME, new MicroTimestampConverter());
-    LogicalConverterRegistry.register(Time.SCHEMA_NAME, new TimeConverter());
-    LogicalConverterRegistry.register(ZonedTimestamp.SCHEMA_NAME, new ZonedTimestampConverter());
-    LogicalConverterRegistry.register(Timestamp.SCHEMA_NAME, new TimestampConverter(config.getShouldConvertDebeziumTimestampToInteger()));
-    LogicalConverterRegistry.register(VariableScaleDecimal.LOGICAL_NAME, new VariableScaleDecimalConverter(config.getVariableScaleDecimalHandlingMode()));
+    LogicalConverterRegistry.registerIfAbsent(Date.SCHEMA_NAME, new DateConverter());
+    LogicalConverterRegistry.registerIfAbsent(MicroTime.SCHEMA_NAME, new MicroTimeConverter());
+    LogicalConverterRegistry.registerIfAbsent(MicroTimestamp.SCHEMA_NAME, new MicroTimestampConverter());
+    LogicalConverterRegistry.registerIfAbsent(Time.SCHEMA_NAME, new TimeConverter());
+    LogicalConverterRegistry.registerIfAbsent(ZonedTimestamp.SCHEMA_NAME, new ZonedTimestampConverter());
+    LogicalConverterRegistry.registerIfAbsent(Timestamp.SCHEMA_NAME, new TimestampConverter(config.getShouldConvertDebeziumTimestampToInteger()));
+    LogicalConverterRegistry.registerIfAbsent(VariableScaleDecimal.LOGICAL_NAME, new VariableScaleDecimalConverter(config.getVariableScaleDecimalHandlingMode()));
+  }
+
+  /**
+   * Remove the DebeziumLogicalConverters from the LogicalConverterRegistry.
+   */
+  public static void remove() {
+    LogicalConverterRegistry.unregister(Date.SCHEMA_NAME);
+    LogicalConverterRegistry.unregister(MicroTime.SCHEMA_NAME);
+    LogicalConverterRegistry.unregister(MicroTimestamp.SCHEMA_NAME);
+    LogicalConverterRegistry.unregister(Time.SCHEMA_NAME);
+    LogicalConverterRegistry.unregister(ZonedTimestamp.SCHEMA_NAME);
+    LogicalConverterRegistry.unregister(Timestamp.SCHEMA_NAME);
+    LogicalConverterRegistry.unregister(VariableScaleDecimal.LOGICAL_NAME);
   }
 
   private DebeziumLogicalConverters() {
diff --git a/kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/convert/logicaltype/KafkaLogicalConverters.java b/kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/convert/logicaltype/KafkaLogicalConverters.java
@@ -52,11 +52,23 @@ public class KafkaLogicalConverters {
   private static final int MAX_NUMERIC_PRECISION = 38;
   private static final int MAX_NUMERIC_SCALE = 9;
 
+  /**
+   * Initialize the LogicalConverterRegistry with the KafkaLogicalConverters.
+   *
+   * @param config the configuration to use.
+   */
   public static void initialize(final BigQuerySinkConfig config) {
-    LogicalConverterRegistry.register(Date.LOGICAL_NAME, new DateConverter());
-    LogicalConverterRegistry.register(Decimal.LOGICAL_NAME, new DecimalConverter(config.getDecimalHandlingMode()));
-    LogicalConverterRegistry.register(Timestamp.LOGICAL_NAME, new TimestampConverter());
-    LogicalConverterRegistry.register(Time.LOGICAL_NAME, new TimeConverter());
+    LogicalConverterRegistry.registerIfAbsent(Date.LOGICAL_NAME, new DateConverter());
+    LogicalConverterRegistry.registerIfAbsent(Decimal.LOGICAL_NAME, new DecimalConverter(config.getDecimalHandlingMode()));
+    LogicalConverterRegistry.registerIfAbsent(Timestamp.LOGICAL_NAME, new TimestampConverter());
+    LogicalConverterRegistry.registerIfAbsent(Time.LOGICAL_NAME, new TimeConverter());
+  }
+
+  public static void remove() {
+    LogicalConverterRegistry.unregister(Date.LOGICAL_NAME);
+    LogicalConverterRegistry.unregister(Decimal.LOGICAL_NAME);
+    LogicalConverterRegistry.unregister(Timestamp.LOGICAL_NAME);
+    LogicalConverterRegistry.unregister(Time.LOGICAL_NAME);
   }
 
   private KafkaLogicalConverters() {
diff --git a/kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/convert/logicaltype/LogicalConverterRegistry.java b/kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/convert/logicaltype/LogicalConverterRegistry.java
@@ -79,7 +79,7 @@ public static LogicalTypeConverter getConverter(String logicalTypeName) {
    * Determines if a converter is registered with the logical type name.
    *
    * @param typeName the logical type name.
-   * @return }{@code true} if there is a converter registered, {@code false} otherwise.
+   * @return {@code true} if there is a converter registered, {@code false} otherwise.
    */
   public static boolean isRegisteredLogicalType(String typeName) {
     return typeName != null && converterMap.containsKey(typeName);
diff --git a/kcbq-connector/src/test/java/com/wepay/kafka/connect/bigquery/convert/BigQueryRecordConverterTest.java b/kcbq-connector/src/test/java/com/wepay/kafka/connect/bigquery/convert/BigQueryRecordConverterTest.java
@@ -55,6 +55,7 @@
 import org.apache.kafka.connect.data.Struct;
 import org.apache.kafka.connect.data.Timestamp;
 import org.apache.kafka.connect.sink.SinkRecord;
+import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -71,8 +72,17 @@ public class BigQueryRecordConverterTest {
   void resetValues() {
     shouldConvertDouble = true;
     useStorageWriteApiConfig = false;
+    DebeziumLogicalConverters.remove();
+    KafkaLogicalConverters.remove();
   }
 
+  @AfterAll
+  static void cleanUpConverters() {
+    DebeziumLogicalConverters.remove();
+    KafkaLogicalConverters.remove();
+  }
+
+
   private static SinkRecord spoofSinkRecord(Schema schema, Object struct, boolean isKey) {
     if (isKey) {
       return new SinkRecord(null, 0, schema, struct, null, null, 0);
diff --git a/kcbq-connector/src/test/java/com/wepay/kafka/connect/bigquery/convert/BigQuerySchemaConverterTest.java b/kcbq-connector/src/test/java/com/wepay/kafka/connect/bigquery/convert/BigQuerySchemaConverterTest.java
@@ -42,6 +42,7 @@
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.SchemaBuilder;
 import org.apache.kafka.connect.data.Timestamp;
+import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -63,6 +64,14 @@ public class BigQuerySchemaConverterTest {
   void resetValues() {
     allFieldsNullable = false;
     sanitizeFieldNames = false;
+    DebeziumLogicalConverters.remove();
+    KafkaLogicalConverters.remove();
+  }
+
+  @AfterAll
+  static void cleanUpConverters() {
+    DebeziumLogicalConverters.remove();
+    KafkaLogicalConverters.remove();
   }
 
   private static BigQuerySinkConfig testingConfig(boolean convertDebeziumTimestamp, BigQuerySinkConfig.DecimalHandlingMode varibaleScaleDecimalMode, BigQuerySinkConfig.DecimalHandlingMode decimalMode) {

Original file line number	Diff line number	Diff line change
`@@ -73,8 +73,6 @@ public void start(Map<String, String> properties) {`
`73`	`73`	`logger.trace("connector.start()");`
`74`	`74`	`configProperties = properties;`
`75`	`75`	`config = new BigQuerySinkConfig(properties);`
`76`		`- DebeziumLogicalConverters.initialize(config);`
`77`		`- KafkaLogicalConverters.initialize(config);`
`78`	`76`	`}`
`79`	`77`
`80`	`78`	`@Override`
Original file line number	Diff line number	Diff line change
`@@ -79,7 +79,7 @@ public static LogicalTypeConverter getConverter(String logicalTypeName) {`
`79`	`79`	`* Determines if a converter is registered with the logical type name.`
`80`	`80`	`*`
`81`	`81`	`* @param typeName the logical type name.`
`82`		`- * @return }{@code true} if there is a converter registered, {@code false} otherwise.`
	`82`	`+ * @return {@code true} if there is a converter registered, {@code false} otherwise.`
`83`	`83`	`*/`
`84`	`84`	`public static boolean isRegisteredLogicalType(String typeName) {`
`85`	`85`	`return typeName != null && converterMap.containsKey(typeName);`