make export and stream works together.

Signed-off-by: Haidong <[email protected]>
wanghd89 · Nov 6, 2023 · 04a8ec2 · 04a8ec2
1 parent 2c39825
commit 04a8ec2
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 7 deletions.
diff --git a/.../main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/KafkaConnectSource.java b/.../main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/KafkaConnectSource.java
@@ -34,7 +34,7 @@
 @SuppressWarnings("deprecation")
 public abstract class KafkaConnectSource implements Source<Record<Object>> {
     private static final Logger LOG = LoggerFactory.getLogger(KafkaConnectSource.class);
-    private final ConnectorConfig connectorConfig;
+    public final ConnectorConfig connectorConfig;
     private final String pipelineName;
     private KafkaConnectConfig kafkaConnectConfig;
     private KafkaConnect kafkaConnect;

diff --git a/...s/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/MongoDBSource.java b/...s/src/main/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/MongoDBSource.java
@@ -40,8 +40,6 @@ public class MongoDBSource extends KafkaConnectSource implements UsesSourceCoord
 
     private final PluginMetrics pluginMetrics;
 
-    private final MongoDBConfig mongoDBConfig;
-
     private final AcknowledgementSetManager acknowledgementSetManager;
 
     private MongoDBService mongoDBService;
@@ -60,7 +58,6 @@ public MongoDBSource(final MongoDBConfig mongoDBConfig,
                          final KafkaConnectConfigSupplier kafkaConnectConfigSupplier) {
         super(mongoDBConfig, pluginMetrics, pipelineDescription, kafkaClusterConfigSupplier, kafkaConnectConfigSupplier);
         this.pluginMetrics = pluginMetrics;
-        this.mongoDBConfig = mongoDBConfig;
         this.acknowledgementSetManager = acknowledgementSetManager;
         this.awsCredentialsSupplier = awsCredentialsSupplier;
         this.byteDecoder = new JsonDecoder();
@@ -71,7 +68,7 @@ public void start(Buffer<Record<Object>> buffer) {
         super.start(buffer);
         if (shouldStartInitialLoad()) {
             LOG.info("Starting initial load");
-            this.mongoDBService = MongoDBService.create(mongoDBConfig, sourceCoordinator, buffer, acknowledgementSetManager, pluginMetrics);
+            this.mongoDBService = MongoDBService.create((MongoDBConfig) this.connectorConfig, sourceCoordinator, buffer, acknowledgementSetManager, pluginMetrics);
             this.mongoDBService.start();
         }
     }
@@ -103,11 +100,12 @@ public ByteDecoder getDecoder() {
 
     @Override
     public boolean shouldStartKafkaConnect() {
-        return false;
-//        return mongoDBConfig.getSnapshotMode().equals("export_stream") || mongoDBConfig.getSnapshotMode().equals("stream");
+        final MongoDBConfig mongoDBConfig = (MongoDBConfig) this.connectorConfig;
+        return mongoDBConfig.getIngestionMode().equals("export_stream") || mongoDBConfig.getIngestionMode().equals("stream");
     }
 
     private boolean shouldStartInitialLoad() {
+        final MongoDBConfig mongoDBConfig = (MongoDBConfig) this.connectorConfig;
         return mongoDBConfig.getIngestionMode().equals("export_stream") || mongoDBConfig.getIngestionMode().equals("export");
     }
 }
diff --git a/...c/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/MongoDBSourceTest.java b/...c/test/java/org/opensearch/dataprepper/plugins/kafkaconnect/source/MongoDBSourceTest.java
@@ -5,6 +5,7 @@
 
 package org.opensearch.dataprepper.plugins.kafkaconnect.source;
 
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;
@@ -16,6 +17,7 @@
 import org.opensearch.dataprepper.plugins.kafkaconnect.configuration.MongoDBConfig;
 
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 @ExtendWith(MockitoExtension.class)
@@ -35,6 +37,11 @@ public class MongoDBSourceTest {
     @Mock
     private AcknowledgementSetManager acknowledgementSetManager;
 
+    @BeforeEach
+    void setup() {
+        mongoDBConfig = mock(MongoDBConfig.class);
+        when(mongoDBConfig.getIngestionMode()).thenReturn("export_stream");
+    }
 
     @Test
     void testConstructorValidations() {