Get shard iterator at sequence number for last shard iterators on ending sequence numbers (#6251)

Signed-off-by: Taylor Gray <tylgry@amazon.com>

Author: graytaylor0

data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/coordination/state/LeaderProgressState.java

@@ -27,4 +27,5 @@ public List<String> getStreamArns() {
     public void setStreamArns(List<String> streamArns) {
         this.streamArns = streamArns;
     }
+
 }

data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/leader/LeaderScheduler.java

@@ -362,5 +362,4 @@ private void createExportPartition(String tableArn, Instant exportTime, String b
         ExportPartition exportPartition = new ExportPartition(tableArn, exportTime, Optional.of(exportProgressState));
         coordinator.createPartition(exportPartition);
     }
-
 }

data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactory.java

@@ -86,11 +86,11 @@ public Runnable createConsumer(final StreamPartition streamPartition,
             // If ending sequence number is present, get the shardIterator for last record
             String endingSequenceNumber = progressState.get().getEndingSequenceNumber();
             if (endingSequenceNumber != null && !endingSequenceNumber.isEmpty()) {
-                lastShardIterator = getShardIterator(streamPartition.getStreamArn(), streamPartition.getShardId(), endingSequenceNumber);
+                lastShardIterator = getShardIterator(streamPartition.getStreamArn(), streamPartition.getShardId(), endingSequenceNumber, ShardIteratorType.AT_SEQUENCE_NUMBER);
             }
         }
 
-        String shardIterator = getShardIterator(streamPartition.getStreamArn(), streamPartition.getShardId(), sequenceNumber);
+        String shardIterator = getShardIterator(streamPartition.getStreamArn(), streamPartition.getShardId(), sequenceNumber, ShardIteratorType.AFTER_SEQUENCE_NUMBER);
         if (shardIterator == null) {
             LOG.error("Failed to start consuming shard '{}'. Unable to get a shard iterator for this shard, this shard may have expired", streamPartition.getShardId());
             return null;
@@ -136,7 +136,7 @@ private TableInfo getTableInfo(String tableArn) {
      * @param sequenceNumber The last Sequence Number processed if any
      * @return A shard iterator.
      */
-    public String getShardIterator(String streamArn, String shardId, String sequenceNumber) {
+    public String getShardIterator(String streamArn, String shardId, String sequenceNumber, ShardIteratorType shardIteratorType) {
         LOG.debug("Get Initial Shard Iter for {}", shardId);
         GetShardIteratorRequest getShardIteratorRequest;
 
@@ -145,7 +145,7 @@ public String getShardIterator(String streamArn, String shardId, String sequence
             getShardIteratorRequest = GetShardIteratorRequest.builder()
                     .shardId(shardId)
                     .streamArn(streamArn)
-                    .shardIteratorType(ShardIteratorType.AFTER_SEQUENCE_NUMBER)
+                    .shardIteratorType(shardIteratorType)
                     .sequenceNumber(sequenceNumber)
                     .build();
         } else {

data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/leader/LeaderSchedulerTest.java

@@ -189,8 +189,14 @@ void test_shardDiscovery_should_create_children_partitions() throws InterruptedE
         leaderPartition = new LeaderPartition();
         leaderPartition.getProgressState().get().setInitialized(true);
         leaderPartition.getProgressState().get().setStreamArns(List.of(streamArn));
+
+        final String expectedEndingSequenceNumber = UUID.randomUUID().toString();
+        when(shardManager.getEndingSequenceNumber(anyString())).thenReturn(expectedEndingSequenceNumber);
         given(coordinator.acquireAvailablePartition(LeaderPartition.PARTITION_TYPE)).willReturn(Optional.of(leaderPartition));
 
+        final ArgumentCaptor<StreamPartition> createdPartitionCaptor = ArgumentCaptor.forClass(StreamPartition.class);
+        when(coordinator.createPartition(createdPartitionCaptor.capture())).thenReturn(true);
+
         ExecutorService executorService = Executors.newSingleThreadExecutor();
         executorService.submit(() -> leaderScheduler.run());
 
@@ -207,6 +213,10 @@ void test_shardDiscovery_should_create_children_partitions() throws InterruptedE
 
         verify(coordinator).saveProgressStateForPartition(eq(leaderPartition), any(Duration.class));
 
+        final List<StreamPartition> createdPartitions = createdPartitionCaptor.getAllValues();
+        createdPartitions.forEach(streamPartition -> {
+            assertThat(streamPartition.getProgressState().get().getEndingSequenceNumber(), equalTo(expectedEndingSequenceNumber));
+        });
     }
 
     @Test

data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/stream/ShardConsumerFactoryTest.java

@@ -32,6 +32,7 @@
 
 import java.time.Duration;
 import java.time.Instant;
+import java.util.List;
 import java.util.Optional;
 import java.util.UUID;
 
@@ -193,7 +194,16 @@ public void test_create_shardConsumer_for_closedShards() {
         Runnable consumer = consumerFactory.createConsumer(streamPartition, Duration.ofMinutes(1), shardAcknowledgementManager);
         assertThat(consumer, notNullValue());
         // Should get iterators twice
-        verify(dynamoDbStreamsClient, times(2)).getShardIterator(any(GetShardIteratorRequest.class));
+
+        final ArgumentCaptor<GetShardIteratorRequest> captor = ArgumentCaptor.forClass(GetShardIteratorRequest.class);
+        verify(dynamoDbStreamsClient, times(2)).getShardIterator(captor.capture());
+        final List<GetShardIteratorRequest> getShardIteratorRequests = captor.getAllValues();
+
+        final GetShardIteratorRequest lastShardIteratorRequest = getShardIteratorRequests.get(0);
+        assertThat(lastShardIteratorRequest.shardIteratorType(), equalTo(ShardIteratorType.AT_SEQUENCE_NUMBER));
+
+        final GetShardIteratorRequest getShardIteratorRequest = getShardIteratorRequests.get(1);
+        assertThat(getShardIteratorRequest.shardIteratorType(), equalTo(ShardIteratorType.TRIM_HORIZON));
 
         verify(streamApiInvocations, times(2)).increment();