Kafka connector was throwing null pointer exception when entire record is null, issue fix. (#2575)

Co-authored-by: ejeffrli <[email protected]>
Co-authored-by: ejeffrli <[email protected]>

Author: VenkatasivareddyTR

athena-kafka/src/main/java/com/amazonaws/athena/connectors/kafka/KafkaRecordHandler.java

@@ -262,6 +262,10 @@ private void execute(
             ConsumerRecord<String, TopicResultSet> record)
     {
         final boolean[] isExecuted = {false};
+        if (record == null || record.value() == null) {
+            LOGGER.warn("[NullRecord] {} Received a null record or record value, offset: {}", splitParameters, record != null ? record.offset() : "unknown");
+            return;
+        }
         spiller.writeRows((Block block, int rowNum) -> {
             for (KafkaField field : record.value().getFields()) {
                 boolean isMatched = block.offerValue(field.getName(), rowNum, field.getValue());
@@ -345,6 +349,10 @@ private void avroExecute(
             SplitParameters splitParameters,
             ConsumerRecord<String, GenericRecord> record)
     {
+        if (record == null || record.value() == null) {
+            LOGGER.warn("[NullRecord] {} Received a null record or record value, offset: {}", splitParameters, record != null ? record.offset() : "unknown");
+            return; // Skip processing this record
+        }
         spiller.writeRows((Block block, int rowNum) -> {
             for (Schema.Field next : record.value().getSchema().getFields()) {
                 boolean isMatched = block.offerValue(next.name(), rowNum, record.value().get(next.name()));
@@ -427,6 +435,10 @@ private void protobufExecute(
             SplitParameters splitParameters,
             ConsumerRecord<String, DynamicMessage> record)
     {
+        if (record == null || record.value() == null) {
+            LOGGER.warn("[NullRecord] {} Received a null record or record value, offset: {}", splitParameters, record != null ? record.offset() : "unknown");
+            return; // Skip processing this record
+        }
         spiller.writeRows((Block block, int rowNum) -> {
             for (Descriptors.FieldDescriptor next : record.value().getAllFields().keySet()) {
                 boolean isMatched = block.offerValue(next.getName(), rowNum, record.value().getField(next));

athena-kafka/src/test/java/com/amazonaws/athena/connectors/kafka/KafkaRecordHandlerTest.java

@@ -122,21 +122,27 @@ public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
 
         consumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
-        ConsumerRecord<String, TopicResultSet> record1 = createConsumerRecord("myTopic", 0, "k1", createTopicResultSet("myTopic"));
-        ConsumerRecord<String, TopicResultSet> record2 = createConsumerRecord("myTopic", 0, "k2", createTopicResultSet("myTopic"));
+        ConsumerRecord<String, TopicResultSet> record1 = createConsumerRecord("myTopic", 0, 0, "k1", createTopicResultSet("myTopic"));
+        ConsumerRecord<String, TopicResultSet> record2 = createConsumerRecord("myTopic", 0, 1, "k2", createTopicResultSet("myTopic"));
+        ConsumerRecord<String, TopicResultSet> nullValueRecord = createConsumerRecord("myTopic", 0, 2, "k3", null);
         consumer.schedulePollTask(() -> {
             consumer.addRecord(record1);
             consumer.addRecord(record2);
+            consumer.addRecord(nullValueRecord);
         });
         avroConsumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
-        ConsumerRecord<String, GenericRecord> avroRecord = createAvroConsumerRecord("greetings", 0 , "k1", createGenericRecord("greetings"));
+        ConsumerRecord<String, GenericRecord> avroRecord = createAvroConsumerRecord("greetings", 0 , 0, "k1", createGenericRecord("greetings"));
+        ConsumerRecord<String, GenericRecord> avroNullValueRecord = createAvroConsumerRecord("greetings", 0 , 1, "k2", null);
         avroConsumer.schedulePollTask(() -> {
             avroConsumer.addRecord(avroRecord);
+            avroConsumer.addRecord(avroNullValueRecord);
         });
         protobufConsumer = new MockConsumer<>(OffsetResetStrategy.EARLIEST);
-        ConsumerRecord<String, DynamicMessage> protobufRecord = createProtobufConsumerRecord("protobuftest", 0, "k1", createDynamicRecord());
+        ConsumerRecord<String, DynamicMessage> protobufRecord = createProtobufConsumerRecord("protobuftest", 0, 0, "k1", createDynamicRecord());
+        ConsumerRecord<String, DynamicMessage> protobufNullValueRecord = createProtobufConsumerRecord("protobuftest", 0, 1, "k2", null);
         protobufConsumer.schedulePollTask(() -> {
             protobufConsumer.addRecord(protobufRecord);
+            protobufConsumer.addRecord(protobufNullValueRecord);
         });
         spillConfig = SpillConfig.newBuilder()
                 .withEncryptionKey(encryptionKey)
@@ -173,7 +179,7 @@ public void testForConsumeDataFromTopic() throws Exception {
         offsets.put(new TopicPartition("myTopic", 0), 1L);
         consumer.updateEndOffsets(offsets);
 
-        SplitParameters splitParameters = new SplitParameters("myTopic", 0, 0, 1);
+        SplitParameters splitParameters = new SplitParameters("myTopic", 0, 0, 2);
         Schema schema = createSchema(createCsvTopicSchema());
 
         mockedKafkaUtils.when(() -> KafkaUtils.getKafkaConsumer(schema, com.google.common.collect.ImmutableMap.of())).thenReturn(consumer);
@@ -351,16 +357,16 @@ private ReadRecordsRequest createReadRecordsRequest(Schema schema) {
                 0);
     }
 
-    private ConsumerRecord<String, TopicResultSet> createConsumerRecord(String topic, int partition, String key, TopicResultSet data) throws Exception {
-        return new ConsumerRecord<>(topic, partition, 0, key, data);
+    private ConsumerRecord<String, TopicResultSet> createConsumerRecord(String topic, int partition, long offset, String key, TopicResultSet data) throws Exception {
+        return new ConsumerRecord<>(topic, partition, offset, key, data);
     }
 
-    private ConsumerRecord<String, GenericRecord> createAvroConsumerRecord(String topic, int partition, String key, GenericRecord data) throws Exception {
-        return new ConsumerRecord<>(topic, partition, 0, key, data);
+    private ConsumerRecord<String, GenericRecord> createAvroConsumerRecord(String topic, int partition, long offset, String key, GenericRecord data) throws Exception {
+        return new ConsumerRecord<>(topic, partition, offset, key, data);
     }
 
-    private ConsumerRecord<String, DynamicMessage> createProtobufConsumerRecord(String topic, int partition, String key, DynamicMessage data) throws Exception {
-        return new ConsumerRecord<>(topic, partition, 0, key, data);
+    private ConsumerRecord<String, DynamicMessage> createProtobufConsumerRecord(String topic, int partition, long offset, String key, DynamicMessage data) throws Exception {
+        return new ConsumerRecord<>(topic, partition, offset, key, data);
     }
 
     private TopicResultSet createTopicResultSet(String topic) {