Merge branch 'main' into storage-write-api-enable-partition-decorator

Author: podgaietska

.github/workflows/nightly.yml

@@ -28,6 +28,12 @@ on:
   workflow_dispatch:
   schedule: ## run GMT 1:17 hours
     - cron: '17 1 * * *'
+  workflow_call:
+    secrets:
+      GCP_CREDENTIALS:
+      KCBQ_TEST_PROJECT:
+      KCBQ_TEST_DATASET:
+      KCBQ_TEST_BUCKET:
 
 permissions:
   contents: write
@@ -44,3 +50,8 @@ permissions:
 jobs:
   call-workflow-2-in-local-repo:
     uses: ./.github/workflows/manual.yml
+    secrets:
+      GCP_CREDENTIALS:
+      KCBQ_TEST_PROJECT:
+      KCBQ_TEST_DATASET:
+      KCBQ_TEST_BUCKET:

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

@@ -26,20 +26,58 @@
 
 import com.google.cloud.bigquery.Field;
 import com.google.cloud.bigquery.LegacySQLTypeName;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
 import java.util.HashMap;
 import java.util.Map;
 import org.apache.kafka.connect.sink.SinkRecord;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * Helper class to construct schema and record for Kafka Data Field.
  */
 public class KafkaDataBuilder {
 
+  private static final Logger logger = LoggerFactory.getLogger(KafkaDataBuilder.class);
+
   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";
 
+  // This is a marker variable for methods necessary to keep original sink record metadata.
+  // These methods in SinkRecord class are available only since Kafka Connect API version 3.6.
+  private static final boolean KAFKA_CONNECT_API_POST_3_6;
+
+  static {
+    boolean kafkaConnectApiPost36;
+    try {
+      MethodHandles.lookup().findVirtual(
+          SinkRecord.class,
+          "originalTopic",
+          MethodType.methodType(String.class)
+      );
+      MethodHandles.lookup().findVirtual(
+          SinkRecord.class,
+          "originalKafkaPartition",
+          MethodType.methodType(Integer.class)
+      );
+      MethodHandles.lookup().findVirtual(
+          SinkRecord.class,
+          "originalKafkaOffset",
+          MethodType.methodType(long.class)
+      );
+      kafkaConnectApiPost36 = true;
+    } catch (NoSuchMethodException | IllegalAccessException e) {
+      logger.warn("This connector cannot retain original topic/partition/offset fields in SinkRecord. "
+                + "If these fields are mutated in upstream SMTs, they will be lost. "
+                + "Upgrade to Kafka Connect 3.6 to provision reliable metadata into resulting table.", e);
+      kafkaConnectApiPost36 = false;
+    }
+    KAFKA_CONNECT_API_POST_3_6 = kafkaConnectApiPost36;
+  }
+
   /**
    * Construct schema for Kafka Data Field
    *
@@ -60,6 +98,30 @@ public static Field buildKafkaDataField(String kafkaDataFieldName) {
         .setMode(com.google.cloud.bigquery.Field.Mode.NULLABLE).build();
   }
 
+  private static String tryGetOriginalTopic(SinkRecord kafkaConnectRecord) {
+    if (KAFKA_CONNECT_API_POST_3_6) {
+      return kafkaConnectRecord.originalTopic();
+    } else {
+      return kafkaConnectRecord.topic();
+    }
+  }
+
+  private static Integer tryGetOriginalKafkaPartition(SinkRecord kafkaConnectRecord) {
+    if (KAFKA_CONNECT_API_POST_3_6) {
+      return kafkaConnectRecord.originalKafkaPartition();
+    } else {
+      return kafkaConnectRecord.kafkaPartition();
+    }
+  }
+
+  private static long tryGetOriginalKafkaOffset(SinkRecord kafkaConnectRecord) {
+    if (KAFKA_CONNECT_API_POST_3_6) {
+      return kafkaConnectRecord.originalKafkaOffset();
+    } else {
+      return kafkaConnectRecord.kafkaOffset();
+    }
+  }
+
   /**
    * Construct a map of Kafka Data record
    *
@@ -68,9 +130,9 @@ public static Field buildKafkaDataField(String kafkaDataFieldName) {
    */
   public static Map<String, Object> buildKafkaDataRecord(SinkRecord kafkaConnectRecord) {
     HashMap<String, Object> kafkaData = new HashMap<>();
-    kafkaData.put(KAFKA_DATA_TOPIC_FIELD_NAME, kafkaConnectRecord.topic());
-    kafkaData.put(KAFKA_DATA_PARTITION_FIELD_NAME, kafkaConnectRecord.kafkaPartition());
-    kafkaData.put(KAFKA_DATA_OFFSET_FIELD_NAME, kafkaConnectRecord.kafkaOffset());
+    kafkaData.put(KAFKA_DATA_TOPIC_FIELD_NAME, tryGetOriginalTopic(kafkaConnectRecord));
+    kafkaData.put(KAFKA_DATA_PARTITION_FIELD_NAME, tryGetOriginalKafkaPartition(kafkaConnectRecord));
+    kafkaData.put(KAFKA_DATA_OFFSET_FIELD_NAME, tryGetOriginalKafkaOffset(kafkaConnectRecord));
     kafkaData.put(KAFKA_DATA_INSERT_TIME_FIELD_NAME, System.currentTimeMillis() / 1000.0);
     return kafkaData;
   }

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

@@ -45,6 +45,10 @@ public class KafkaDataConverterTest {
   private static final String kafkaDataTopicValue = "testTopic";
   private static final int kafkaDataPartitionValue = 101;
   private static final long kafkaDataOffsetValue = 1337;
+  private static final String kafkaDataMutatedTopicValue = "mutatedTopic";
+  private static final int kafkaDataMutatedPartitionValue = 201;
+  // In 3.6.1, there is no direct way to modify offset via newRecord(), even if SinkRecord itself supports it
+  private static final long kafkaDataMutatedOffsetValue = 456;
   Map<String, Object> expectedKafkaDataFields = new HashMap<>();
 
   @BeforeEach
@@ -67,6 +71,33 @@ public void testBuildKafkaDataRecord() {
     assertEquals(expectedKafkaDataFields, actualKafkaDataFields);
   }
 
+  @Test
+  public void testBuildKafkaDataRecordOnMutatedMetadata() {
+    SinkRecord record = new SinkRecord(
+            kafkaDataTopicValue,
+            kafkaDataPartitionValue,
+            null,
+            null,
+            null,
+            null,
+            kafkaDataOffsetValue
+    );
+    SinkRecord mutatedRecord = record.newRecord(
+            kafkaDataMutatedTopicValue,
+            kafkaDataMutatedPartitionValue,
+            null,
+            null,
+            null,
+            null,
+            null
+    );
+
+    Map<String, Object> actualKafkaDataFields = KafkaDataBuilder.buildKafkaDataRecord(mutatedRecord);
+    actualKafkaDataFields.remove(kafkaDataInsertTimeName);
+
+    assertEquals(expectedKafkaDataFields, actualKafkaDataFields);
+  }
+
   @Test
   public void testBuildKafkaDataRecordStorageWriteApi() {
     SinkRecord record = new SinkRecord(kafkaDataTopicValue, kafkaDataPartitionValue, null, null, null, null, kafkaDataOffsetValue);