Update Kafka destination to use outputRecordCollector to properly store state (#15287)

* Added tracking state to destination kafka

* Bump version

* auto-bump connector version [ci skip]

Co-authored-by: Octavia Squidington III <[email protected]>

Author: suhomud

airbyte-config/init/src/main/resources/seed/destination_definitions.yaml

@@ -133,7 +133,7 @@
 - name: Kafka
   destinationDefinitionId: 9f760101-60ae-462f-9ee6-b7a9dafd454d
   dockerRepository: airbyte/destination-kafka
-  dockerImageTag: 0.1.9
+  dockerImageTag: 0.1.10
   documentationUrl: https://docs.airbyte.io/integrations/destinations/kafka
   icon: kafka.svg
   releaseStage: alpha

airbyte-config/init/src/main/resources/seed/destination_specs.yaml

@@ -2005,7 +2005,7 @@
     supportsDBT: false
     supported_destination_sync_modes:
     - "append"
-- dockerImage: "airbyte/destination-kafka:0.1.9"
+- dockerImage: "airbyte/destination-kafka:0.1.10"
   spec:
     documentationUrl: "https://docs.airbyte.io/integrations/destinations/kafka"
     connectionSpecification:

airbyte-integrations/connectors/destination-kafka/Dockerfile

@@ -16,5 +16,5 @@ ENV APPLICATION destination-kafka
 
 COPY --from=build /airbyte /airbyte
 
-LABEL io.airbyte.version=0.1.9
+LABEL io.airbyte.version=0.1.10
 LABEL io.airbyte.name=airbyte/destination-kafka

airbyte-integrations/connectors/destination-kafka/build.gradle

@@ -17,6 +17,8 @@ dependencies {
     implementation 'org.apache.kafka:kafka-clients:2.8.0'
     implementation 'org.apache.kafka:connect-json:2.8.0'
 
+    testImplementation project(':airbyte-integrations:bases:standard-destination-test')
+
     integrationTestJavaImplementation project(':airbyte-integrations:bases:standard-destination-test')
     integrationTestJavaImplementation project(':airbyte-integrations:connectors:destination-kafka')
     integrationTestJavaImplementation libs.connectors.testcontainers.kafka

airbyte-integrations/connectors/destination-kafka/src/main/java/io/airbyte/integrations/destination/kafka/KafkaRecordConsumer.java

@@ -37,8 +37,6 @@ public class KafkaRecordConsumer extends FailureTrackingAirbyteMessageConsumer {
   private final Consumer<AirbyteMessage> outputRecordCollector;
   private final NamingConventionTransformer nameTransformer;
 
-  private AirbyteMessage lastStateMessage = null;
-
   public KafkaRecordConsumer(final KafkaDestinationConfig kafkaDestinationConfig,
                              final ConfiguredAirbyteCatalog catalog,
                              final Consumer<AirbyteMessage> outputRecordCollector,
@@ -60,7 +58,7 @@ protected void startTracked() {
   @Override
   protected void acceptTracked(final AirbyteMessage airbyteMessage) {
     if (airbyteMessage.getType() == AirbyteMessage.Type.STATE) {
-      lastStateMessage = airbyteMessage;
+      outputRecordCollector.accept(airbyteMessage);
     } else if (airbyteMessage.getType() == AirbyteMessage.Type.RECORD) {
       final AirbyteRecordMessage recordMessage = airbyteMessage.getRecord();
 
@@ -98,15 +96,13 @@ private void sendRecord(final ProducerRecord<String, JsonNode> record) {
     });
     if (sync) {
       producer.flush();
-      outputRecordCollector.accept(lastStateMessage);
     }
   }
 
   @Override
   protected void close(final boolean hasFailed) {
     producer.flush();
     producer.close();
-    outputRecordCollector.accept(lastStateMessage);
   }
 
 }

airbyte-integrations/connectors/destination-kafka/src/test/java/io/airbyte/integrations/destination/kafka/KafkaRecordConsumerTest.java

@@ -6,7 +6,6 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.mockito.Mockito.mock;
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -15,7 +14,9 @@
 import io.airbyte.commons.jackson.MoreMappers;
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.integrations.base.AirbyteStreamNameNamespacePair;
+import io.airbyte.integrations.base.FailureTrackingAirbyteMessageConsumer;
 import io.airbyte.integrations.destination.StandardNameTransformer;
+import io.airbyte.integrations.standardtest.destination.PerStreamStateMessageTest;
 import io.airbyte.protocol.models.AirbyteMessage;
 import io.airbyte.protocol.models.AirbyteRecordMessage;
 import io.airbyte.protocol.models.AirbyteStateMessage;
@@ -30,16 +31,21 @@
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.DisplayName;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.api.extension.ExtensionContext;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.ArgumentsProvider;
 import org.junit.jupiter.params.provider.ArgumentsSource;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
 
 @DisplayName("KafkaRecordConsumer")
-public class KafkaRecordConsumerTest {
+@ExtendWith(MockitoExtension.class)
+public class KafkaRecordConsumerTest extends PerStreamStateMessageTest {
 
   private static final ObjectMapper mapper = MoreMappers.initMapper();
   private static final String TOPIC_NAME = "test.topic";
@@ -53,16 +59,27 @@ public class KafkaRecordConsumerTest {
           Field.of("id", JsonSchemaType.NUMBER),
           Field.of("name", JsonSchemaType.STRING))));
 
+  @Mock
+  private Consumer<AirbyteMessage> outputRecordCollector;
+
+  private KafkaRecordConsumer consumer;
+
   private static final StandardNameTransformer NAMING_RESOLVER = new StandardNameTransformer();
 
+  @BeforeEach
+  public void init() {
+    final KafkaDestinationConfig config = KafkaDestinationConfig.getKafkaDestinationConfig(getConfig(TOPIC_NAME));
+    consumer = new KafkaRecordConsumer(config, CATALOG, outputRecordCollector, NAMING_RESOLVER);
+  }
+
   @ParameterizedTest
   @ArgumentsSource(TopicMapArgumentsProvider.class)
   @SuppressWarnings("unchecked")
   public void testBuildTopicMap(final String topicPattern, final String expectedTopic) {
     final KafkaDestinationConfig config = KafkaDestinationConfig.getKafkaDestinationConfig(getConfig(topicPattern));
-    final KafkaRecordConsumer recordConsumer = new KafkaRecordConsumer(config, CATALOG, mock(Consumer.class), NAMING_RESOLVER);
+    consumer = new KafkaRecordConsumer(config, CATALOG, outputRecordCollector, NAMING_RESOLVER);
 
-    final Map<AirbyteStreamNameNamespacePair, String> topicMap = recordConsumer.buildTopicMap();
+    final Map<AirbyteStreamNameNamespacePair, String> topicMap = consumer.buildTopicMap();
     assertEquals(1, topicMap.size());
 
     final AirbyteStreamNameNamespacePair streamNameNamespacePair = new AirbyteStreamNameNamespacePair(STREAM_NAME, SCHEMA_NAME);
@@ -72,8 +89,6 @@ public void testBuildTopicMap(final String topicPattern, final String expectedTo
   @Test
   @SuppressWarnings("unchecked")
   void testCannotConnectToBrokers() throws Exception {
-    final KafkaDestinationConfig config = KafkaDestinationConfig.getKafkaDestinationConfig(getConfig(TOPIC_NAME));
-    final KafkaRecordConsumer consumer = new KafkaRecordConsumer(config, CATALOG, mock(Consumer.class), NAMING_RESOLVER);
     final List<AirbyteMessage> expectedRecords = getNRecords(10);
 
     consumer.start();
@@ -133,6 +148,16 @@ private List<AirbyteMessage> getNRecords(final int n) {
 
   }
 
+  @Override
+  protected Consumer<AirbyteMessage> getMockedConsumer() {
+    return outputRecordCollector;
+  }
+
+  @Override
+  protected FailureTrackingAirbyteMessageConsumer getMessageConsumer() {
+    return consumer;
+  }
+
   public static class TopicMapArgumentsProvider implements ArgumentsProvider {
 
     @Override