diff --git a/src/main/java/io/kestra/plugin/aws/s3/Download.java b/src/main/java/io/kestra/plugin/aws/s3/Download.java index dbb3a5e2..f06058e8 100644 --- a/src/main/java/io/kestra/plugin/aws/s3/Download.java +++ b/src/main/java/io/kestra/plugin/aws/s3/Download.java @@ -126,7 +126,7 @@ public class Download extends AbstractS3Object implements RunnableTask getMaxKeys(); diff --git a/src/main/java/io/kestra/plugin/aws/s3/S3Service.java b/src/main/java/io/kestra/plugin/aws/s3/S3Service.java index a193892a..5b700bd2 100644 --- a/src/main/java/io/kestra/plugin/aws/s3/S3Service.java +++ b/src/main/java/io/kestra/plugin/aws/s3/S3Service.java @@ -7,6 +7,7 @@ import java.util.HashMap; import java.util.List; import java.util.concurrent.ExecutionException; +import java.util.regex.Pattern; import org.apache.commons.io.FilenameUtils; import org.apache.commons.lang3.StringUtils; @@ -32,8 +33,6 @@ import software.amazon.awssdk.transfer.s3.model.DownloadFileRequest; import software.amazon.awssdk.transfer.s3.model.FileDownload; -import static io.kestra.core.utils.Rethrow.throwPredicate; - public class S3Service { public static void initCrt() { @@ -240,11 +239,15 @@ static void performAction( public static List list(RunContext runContext, S3Client client, ListInterface list, AbstractS3Object abstractS3) throws IllegalVariableEvaluationException { int rMaxKeys = runContext.render(list.getMaxKeys()).as(Integer.class).orElse(1000); + if (rMaxKeys < 0) { + throw new IllegalArgumentException("The `maxKeys` property must be 0 or more, got %d.".formatted(rMaxKeys)); + } + if (rMaxKeys == 0) { + return List.of(); + } ListObjectsV2Request.Builder builder = ListObjectsV2Request.builder() - .bucket(runContext.render(list.getBucket()).as(String.class).orElseThrow()) - // per-page size is capped at 1000 by S3, rMaxKeys is the total cap across pages - .maxKeys(Math.min(rMaxKeys, 1000)); + .bucket(runContext.render(list.getBucket()).as(String.class).orElseThrow(() -> new IllegalArgumentException("The `bucket` property is required."))); if (list.getPrefix() != null) { builder.prefix(runContext.render(list.getPrefix()).as(String.class).orElseThrow()); @@ -271,21 +274,26 @@ public static List list(RunContext runContext, S3Client client, ListIn builder.requestPayer(runContext.render(abstractS3.getRequestPayer()).as(String.class).orElseThrow()); } - String regExp = runContext.render(list.getRegexp()).as(String.class).orElse(null); - var filter = runContext.render(list.getFilter()).as(ListInterface.Filter.class).orElseThrow(); + var rPattern = runContext.render(list.getRegexp()).as(String.class).map(Pattern::compile).orElse(null); + var rFilter = runContext.render(list.getFilter()).as(ListInterface.Filter.class).orElse(ListInterface.Filter.BOTH); + + // per-page size is capped at 1000 by S3, rMaxKeys is the total cap across pages, and a regexp or filter can skip most keys so it reads full pages + builder.maxKeys(rPattern == null && rFilter == ListInterface.Filter.BOTH ? Math.min(rMaxKeys, 1000) : 1000); ListObjectsV2Request baseRequest = builder.build(); - List allContents = new ArrayList<>(); + var result = new ArrayList(); ListObjectsV2Request currentRequest = baseRequest; ListObjectsV2Response response; do { response = client.listObjectsV2(currentRequest); - allContents.addAll(response.contents()); - - // maxKeys is a total cap across pages, the per-page size is capped at 1000 by S3. - if (allContents.size() >= rMaxKeys) { - break; + response.contents().stream() + .filter(s3Object -> filter(s3Object, rPattern, rFilter)) + .limit(rMaxKeys - result.size()) + .map(S3Object::of) + .forEach(result::add); + if (result.size() >= rMaxKeys) { + return result; } if (response.isTruncated()) { @@ -299,18 +307,13 @@ public static List list(RunContext runContext, S3Client client, ListIn } } while (response.isTruncated()); - var contents = allContents.size() > rMaxKeys ? allContents.subList(0, rMaxKeys) : allContents; - - return contents.stream() - .filter(throwPredicate(s3Object -> S3Service.filter(s3Object, regExp, filter))) - .map(S3Object::of) - .toList(); + return result; } - private static boolean filter(software.amazon.awssdk.services.s3.model.S3Object object, String regExp, ListInterface.Filter filter) { - return (regExp == null || object.key().matches(regExp)) && - (filter == ListInterface.Filter.BOTH || - (filter == ListInterface.Filter.DIRECTORY && object.key().endsWith("/")) || - (filter == ListInterface.Filter.FILES && !object.key().endsWith("/"))); + private static boolean filter(software.amazon.awssdk.services.s3.model.S3Object object, Pattern pattern, ListInterface.Filter type) { + return (pattern == null || pattern.matcher(object.key()).matches()) && + (type == ListInterface.Filter.BOTH || + (type == ListInterface.Filter.DIRECTORY && object.key().endsWith("/")) || + (type == ListInterface.Filter.FILES && !object.key().endsWith("/"))); } } diff --git a/src/test/java/io/kestra/plugin/aws/s3/S3ServiceListTest.java b/src/test/java/io/kestra/plugin/aws/s3/S3ServiceListTest.java index 4029a81c..cc505dbe 100644 --- a/src/test/java/io/kestra/plugin/aws/s3/S3ServiceListTest.java +++ b/src/test/java/io/kestra/plugin/aws/s3/S3ServiceListTest.java @@ -6,8 +6,8 @@ import io.kestra.core.junit.annotations.KestraTest; import io.kestra.core.models.property.Property; -import io.kestra.core.runners.RunContext; import io.kestra.core.runners.RunContextFactory; +import io.kestra.core.utils.IdUtils; import io.kestra.plugin.aws.s3.models.S3Object; import jakarta.inject.Inject; @@ -16,7 +16,9 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.*; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.Mockito.*; @KestraTest @@ -28,50 +30,138 @@ class S3ServiceListTest { // Verifies the V2 isTruncated -> continuationToken -> next page loop. @Test void list_paginatesAcrossTwoTruncatedPages() throws Exception { - RunContext runContext = runContextFactory.of(); + var mockClient = twoPageClient(List.of("file1.txt", "file2.txt"), List.of("file3.txt")); - var obj1 = software.amazon.awssdk.services.s3.model.S3Object.builder() - .key("file1.txt").size(10L).build(); - var obj2 = software.amazon.awssdk.services.s3.model.S3Object.builder() - .key("file2.txt").size(20L).build(); - var obj3 = software.amazon.awssdk.services.s3.model.S3Object.builder() - .key("file3.txt").size(30L).build(); + var task = listTask() + .maxKeys(Property.ofValue(1000)) + .maxFiles(Property.ofValue(1000)) + .build(); + + var result = S3Service.list(runContextFactory.of(), mockClient, task, task); + + assertThat(result.stream().map(S3Object::getKey).toList(), contains("file1.txt", "file2.txt", "file3.txt")); + + // Verify the second call carried the continuation token. + verify(mockClient, times(2)).listObjectsV2(any(ListObjectsV2Request.class)); + } + + @Test + void shouldContinuePaginationWhenInitialPageDoesNotMatchFilter() throws Exception { + var mockClient = twoPageClient(List.of("archive/old1.txt", "archive/old2.txt"), List.of("data/orders.csv")); + + var task = listTask() + .regexp(Property.ofValue(".*\\.csv")) + .maxKeys(Property.ofValue(2)) + .build(); + + var result = S3Service.list(runContextFactory.of(), mockClient, task, task); + + assertThat(result, hasSize(1)); + assertThat(result.getFirst().getKey(), is("data/orders.csv")); + // a regexp can skip most keys, so pages are requested at the S3 maximum instead of maxKeys + verify(mockClient, times(2)).listObjectsV2(argThat((ListObjectsV2Request req) -> req.maxKeys() == 1000)); + } + + @Test + void shouldStopPaginationWhenMatchingCountReachesMaxKeys() throws Exception { + var mockClient = twoPageClient(List.of("orders-1.csv", "notes.txt"), List.of("orders-2.csv", "orders-3.csv")); + + var task = listTask() + .regexp(Property.ofValue(".*\\.csv")) + .maxKeys(Property.ofValue(2)) + .build(); + + var result = S3Service.list(runContextFactory.of(), mockClient, task, task); + + assertThat(result.stream().map(S3Object::getKey).toList(), contains("orders-1.csv", "orders-2.csv")); + verify(mockClient, times(2)).listObjectsV2(any(ListObjectsV2Request.class)); + } + + @Test + void shouldContinuePaginationWhenInitialPageHasOnlyDirectories() throws Exception { + var mockClient = twoPageClient(List.of("archive/", "data/"), List.of("data/orders.csv")); + + var task = listTask() + .filter(Property.ofValue(ListInterface.Filter.FILES)) + .maxKeys(Property.ofValue(1)) + .build(); + + var result = S3Service.list(runContextFactory.of(), mockClient, task, task); + + assertThat(result, hasSize(1)); + assertThat(result.getFirst().getKey(), is("data/orders.csv")); + verify(mockClient, times(2)).listObjectsV2(argThat((ListObjectsV2Request req) -> req.maxKeys() == 1000)); + } + @Test + void shouldRequestMaxKeysSizedPagesWithoutFilter() throws Exception { + var mockClient = mock(S3Client.class); + when(mockClient.listObjectsV2(any(ListObjectsV2Request.class))) + .thenReturn(ListObjectsV2Response.builder().isTruncated(false).build()); + var task = listTask() + .maxKeys(Property.ofValue(5)) + .build(); + + S3Service.list(runContextFactory.of(), mockClient, task, task); + + verify(mockClient).listObjectsV2(argThat((ListObjectsV2Request req) -> req.maxKeys() == 5)); + } + + @Test + void shouldReturnNothingWithoutCallingS3WhenMaxKeysIsZero() throws Exception { + var mockClient = mock(S3Client.class); + var task = listTask() + .maxKeys(Property.ofValue(0)) + .build(); + + var result = S3Service.list(runContextFactory.of(), mockClient, task, task); + + assertThat(result, empty()); + verify(mockClient, never()).listObjectsV2(any(ListObjectsV2Request.class)); + } + + @Test + void shouldRejectNegativeMaxKeys() { + var mockClient = mock(S3Client.class); + var task = listTask() + .maxKeys(Property.ofValue(-1)) + .build(); + + var exception = assertThrows(IllegalArgumentException.class, () -> S3Service.list(runContextFactory.of(), mockClient, task, task)); + + assertThat(exception.getMessage(), containsString("must be 0 or more")); + verify(mockClient, never()).listObjectsV2(any(ListObjectsV2Request.class)); + } + + private static io.kestra.plugin.aws.s3.List.ListBuilder listTask() { + return io.kestra.plugin.aws.s3.List.builder() + .id("s3-service-list-%s".formatted(IdUtils.create())) + .type(io.kestra.plugin.aws.s3.List.class.getName()) + .bucket(Property.ofValue("my-bucket")); + } + + // Serves firstPage, then secondPage once the request carries the continuation token. + private static S3Client twoPageClient(List firstPage, List secondPage) { var page1 = ListObjectsV2Response.builder() .isTruncated(true) .nextContinuationToken("token-page2") - .contents(obj1, obj2) + .contents(s3Objects(firstPage)) .build(); var page2 = ListObjectsV2Response.builder() .isTruncated(false) - .contents(obj3) + .contents(s3Objects(secondPage)) .build(); - S3Client mockClient = mock(S3Client.class); + var mockClient = mock(S3Client.class); when(mockClient.listObjectsV2(any(ListObjectsV2Request.class))) - .thenAnswer(inv -> - { - ListObjectsV2Request req = inv.getArgument(0); - return "token-page2".equals(req.continuationToken()) ? page2 : page1; - }); - - var task = io.kestra.plugin.aws.s3.List.builder() - .id("s3-service-list-unit-test") - .type(io.kestra.plugin.aws.s3.List.class.getName()) - .bucket(Property.ofValue("my-bucket")) - .maxKeys(Property.ofValue(1000)) - .maxFiles(Property.ofValue(1000)) - .build(); - - List result = S3Service.list(runContext, mockClient, task, task); - - assertThat(result, hasSize(3)); - assertThat(result.get(0).getKey(), is("file1.txt")); - assertThat(result.get(1).getKey(), is("file2.txt")); - assertThat(result.get(2).getKey(), is("file3.txt")); + .thenAnswer(inv -> "token-page2".equals(inv.getArgument(0).continuationToken()) ? page2 : page1); + return mockClient; + } - // Verify the second call carried the continuation token. - verify(mockClient, times(2)).listObjectsV2(any(ListObjectsV2Request.class)); + private static List s3Objects(List keys) { + return keys.stream() + .map(key -> software.amazon.awssdk.services.s3.model.S3Object.builder().key(key).size(10L).build()) + .toList(); } }