Fixed kinesis retrieval config argument passing to KCL scheduler (ope…

…nsearch-project#5272) Signed-off-by: RashmiRam <[email protected]>
san81 · Jan 21, 2025 · 9c61e03 · 9c61e03
1 parent 01c9b4a
commit 9c61e03
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 2 deletions.
diff --git a/...ource/src/main/java/org/opensearch/dataprepper/plugins/kinesis/source/KinesisService.java b/...ource/src/main/java/org/opensearch/dataprepper/plugins/kinesis/source/KinesisService.java
@@ -41,6 +41,7 @@
 import software.amazon.kinesis.exceptions.KinesisClientLibDependencyException;
 import software.amazon.kinesis.exceptions.ThrottlingException;
 import software.amazon.kinesis.processor.ShardRecordProcessorFactory;
+import software.amazon.kinesis.retrieval.RetrievalConfig;
 import software.amazon.kinesis.retrieval.polling.PollingConfig;
 
 import java.time.Duration;
@@ -185,9 +186,10 @@ public Scheduler createScheduler(final Buffer<Record<Event>> buffer) {
                 .tableName(tableName)
                 .namespace(kclMetricsNamespaceName);
 
+        RetrievalConfig retrievalConfig = configsBuilder.retrievalConfig();
         ConsumerStrategy consumerStrategy = kinesisSourceConfig.getConsumerStrategy();
         if (consumerStrategy == ConsumerStrategy.POLLING) {
-            configsBuilder.retrievalConfig().retrievalSpecificConfig(
+            retrievalConfig = configsBuilder.retrievalConfig().retrievalSpecificConfig(
                 new PollingConfig(kinesisClient)
                     .maxRecords(kinesisSourceConfig.getPollingConfig().getMaxPollingRecords())
                     .idleTimeBetweenReadsInMillis(
@@ -203,7 +205,7 @@ public Scheduler createScheduler(final Buffer<Record<Event>> buffer) {
                 configsBuilder.lifecycleConfig(),
                 configsBuilder.metricsConfig(),
                 configsBuilder.processorConfig(),
-                configsBuilder.retrievalConfig()
+                retrievalConfig
         );
     }
 }
diff --git a/...e/src/test/java/org/opensearch/dataprepper/plugins/kinesis/source/KinesisServiceTest.java b/...e/src/test/java/org/opensearch/dataprepper/plugins/kinesis/source/KinesisServiceTest.java
@@ -42,6 +42,7 @@
 import software.amazon.kinesis.common.InitialPositionInStream;
 import software.amazon.kinesis.coordinator.Scheduler;
 import software.amazon.kinesis.metrics.MetricsLevel;
+import software.amazon.kinesis.retrieval.polling.PollingConfig;
 
 import java.time.Duration;
 import java.time.Instant;
@@ -276,6 +277,30 @@ void testCreateSchedulerWithPollingStrategy() {
         verify(workerIdentifierGenerator, times(1)).generate();
     }
 
+    @Test
+    void testCreateSchedulerWithPollingStrategyAndPollingConfig() {
+        when(kinesisSourceConfig.getConsumerStrategy()).thenReturn(ConsumerStrategy.POLLING);
+        when(kinesisSourceConfig.getPollingConfig()).thenReturn(kinesisStreamPollingConfig);
+        KinesisService kinesisService = new KinesisService(kinesisSourceConfig, kinesisClientFactory, pluginMetrics, pluginFactory,
+            pipelineDescription, acknowledgementSetManager, kinesisLeaseConfigSupplier, workerIdentifierGenerator);
+        Scheduler schedulerObjectUnderTest = kinesisService.createScheduler(buffer);
+
+        assertEquals(kinesisService.getApplicationName(), pipelineName);
+        assertNotNull(schedulerObjectUnderTest);
+        assertNotNull(schedulerObjectUnderTest.checkpointConfig());
+        assertNotNull(schedulerObjectUnderTest.leaseManagementConfig());
+        assertSame(schedulerObjectUnderTest.leaseManagementConfig().initialPositionInStream().getInitialPositionInStream(), InitialPositionInStream.TRIM_HORIZON);
+        assertNotNull(schedulerObjectUnderTest.lifecycleConfig());
+        assertNotNull(schedulerObjectUnderTest.metricsConfig());
+        assertSame(schedulerObjectUnderTest.metricsConfig().metricsLevel(), MetricsLevel.DETAILED);
+        assertNotNull(schedulerObjectUnderTest.processorConfig());
+        assertNotNull(schedulerObjectUnderTest.retrievalConfig().retrievalSpecificConfig());
+        assertEquals(((PollingConfig)schedulerObjectUnderTest.retrievalConfig().retrievalSpecificConfig()).maxRecords(), kinesisStreamPollingConfig.getMaxPollingRecords());
+        assertEquals(((PollingConfig)schedulerObjectUnderTest.retrievalConfig().retrievalSpecificConfig()).idleTimeBetweenReadsInMillis(), kinesisStreamPollingConfig.getIdleTimeBetweenReads().toMillis());
+        assertNotNull(schedulerObjectUnderTest.retrievalConfig());
+        verify(workerIdentifierGenerator, times(1)).generate();
+    }
+
     @Test
     void testServiceStartNullBufferThrows() {
         KinesisService kinesisService = createObjectUnderTest();