fix(s3): filter objects during pagination in S3Service.list to avoid premature loop exit

src/main/java/io/kestra/plugin/aws/s3/Download.java

@@ -126,7 +126,7 @@ public class Download extends AbstractS3Object implements RunnableTask<Download.
 
     @Schema(
         title = "Max keys",
-        description = "Maximum keys per list request in multi-file mode; default 1000."
+        description = "Maximum number of files listed in multi-file mode, counted after `regexp`. Default 1000. `maxFiles` caps how many are downloaded."
     )
     @Builder.Default
     @PluginProperty(group = "connection")

src/main/java/io/kestra/plugin/aws/s3/ListInterface.java

@@ -41,7 +41,10 @@ public interface ListInterface {
 
     @Schema(
         title = "Maximum number of objects returned across all paginated S3 calls",
-        description = "Total upper bound on the objects returned. It also sets the per-page size sent to S3, which AWS caps at 1000. The listing pages until this total is reached or the bucket is exhausted. Default 1000."
+        description = """
+            Total upper bound on the objects returned, counted after `regexp` and `filter`. Default 1000, must be 0 or more, and 0 returns no objects.
+            When `regexp` is unset and `filter` is `BOTH`, it also sets the per-page size sent to S3, which AWS caps at 1000.
+            The listing pages until this total is reached or the bucket is exhausted, so a `regexp` or `filter` that matches few keys can read every key under `prefix`. Set `prefix` to narrow it."""
     )
     @PluginProperty(group = "execution")
     Property<Integer> getMaxKeys();

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<S3Object> 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<S3Object> 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<software.amazon.awssdk.services.s3.model.S3Object> allContents = new ArrayList<>();
+        var result = new ArrayList<S3Object>();
 
         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<S3Object> 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("/")));
     }
 }

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<String> firstPage, List<String> 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<S3Object> 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.<ListObjectsV2Request>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<software.amazon.awssdk.services.s3.model.S3Object> s3Objects(List<String> keys) {
+        return keys.stream()
+            .map(key -> software.amazon.awssdk.services.s3.model.S3Object.builder().key(key).size(10L).build())
+            .toList();
     }
 }