From 513331224b7a64fc4580f295524601d1189f80e2 Mon Sep 17 00:00:00 2001 From: Taylor Gray Date: Sat, 30 May 2026 14:04:43 -0500 Subject: [PATCH] Fix bug that would cause ddb stream processing to start before export completed Signed-off-by: Taylor Gray --- .../plugins/source/dynamodb/export/ExportScheduler.java | 2 +- .../source/dynamodb/export/ExportSchedulerTest.java | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java index ed0d67fd41..5fb8b83703 100644 --- a/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java +++ b/data-prepper-plugins/dynamodb-source/src/main/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportScheduler.java @@ -197,7 +197,7 @@ private void createDataFilePartitions(final String exportArn, // Currently, we need to maintain a global state to track the overall progress. // So that we can easily tell if all the export files are loaded - LoadStatus loadStatus = new LoadStatus(totalFiles.get(), 0, totalRecords.get(), 0); + LoadStatus loadStatus = new LoadStatus(dataFileInfo.size(), 0, totalRecords.get(), 0); enhancedSourceCoordinator.createPartition(new GlobalState(exportArn, Optional.of(loadStatus.toMap()))); dataFileInfo.forEach((key, size) -> { diff --git a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportSchedulerTest.java b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportSchedulerTest.java index 383c52bbc7..bf69551ba5 100644 --- a/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportSchedulerTest.java +++ b/data-prepper-plugins/dynamodb-source/src/test/java/org/opensearch/dataprepper/plugins/source/dynamodb/export/ExportSchedulerTest.java @@ -20,6 +20,7 @@ import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.partition.GlobalState; import org.opensearch.dataprepper.plugins.source.dynamodb.coordination.state.ExportProgressState; import org.opensearch.dataprepper.plugins.source.dynamodb.model.ExportSummary; +import org.opensearch.dataprepper.plugins.source.dynamodb.model.LoadStatus; import org.opensearch.dataprepper.plugins.source.dynamodb.utils.DynamoDBSourceAggregateMetrics; import software.amazon.awssdk.services.dynamodb.DynamoDbClient; import software.amazon.awssdk.services.dynamodb.model.DescribeExportRequest; @@ -185,6 +186,12 @@ public void test_run_exportJob_correctly() throws InterruptedException { assertThat(createdPartitions.get(1), instanceOf(DataFilePartition.class)); assertThat(createdPartitions.get(2), instanceOf(DataFilePartition.class)); + // Verify the GlobalState LoadStatus has accurate totalFiles count + GlobalState globalState = (GlobalState) createdPartitions.get(0); + LoadStatus loadStatus = LoadStatus.fromMap(globalState.getProgressState().get()); + assertThat(loadStatus.getTotalFiles(), equalTo(2)); + assertThat(loadStatus.getLoadedFiles(), equalTo(0)); + // Complete the export partition verify(coordinator).completePartition(any(EnhancedSourcePartition.class)); verify(exportJobSuccess).increment();