fix(core): implement allByPrefix

kestra-io · Dec 15, 2023 · 97b7af5 · 97b7af5
1 parent a65317e
commit 97b7af5
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/src/main/java/io/kestra/storage/s3/S3Storage.java b/src/main/java/io/kestra/storage/s3/S3Storage.java
@@ -66,11 +66,9 @@ public InputStream get(String tenantId, URI uri) throws IOException {
     }
 
     @Override
-    public List<URI> filesByPrefix(String tenantId, URI prefix) {
+    public List<URI> allByPrefix(String tenantId, URI prefix, boolean includeDirectories) {
         String path = getPath(tenantId, prefix);
-        return keysForPrefix(path, true)
-            // keep only files
-            .filter(key -> !key.endsWith("/"))
+        return keysForPrefix(path, true, includeDirectories)
             .map(key -> URI.create("kestra://" + prefix.getPath() + key.substring(path.length())))
             .toList();
     }
@@ -102,7 +100,7 @@ public List<FileAttributes> list(String tenantId, URI uri) throws IOException {
         String path = getPath(tenantId, uri);
         String prefix = path.endsWith("/") ? path : path + "/";
         try {
-            List<FileAttributes> list = keysForPrefix(prefix, false)
+            List<FileAttributes> list = keysForPrefix(prefix, false, true)
                 .map(throwFunction(this::getFileAttributes))
                 .toList();
             if (list.isEmpty()) {
@@ -117,7 +115,7 @@ public List<FileAttributes> list(String tenantId, URI uri) throws IOException {
         }
     }
 
-    private Stream<String> keysForPrefix(String prefix, boolean recursive) {
+    private Stream<String> keysForPrefix(String prefix, boolean recursive, boolean includeDirectories) {
         ListObjectsV2Request request = ListObjectsV2Request.builder()
             .bucket(s3Config.getBucket())
             .prefix(prefix)
@@ -130,7 +128,9 @@ private Stream<String> keysForPrefix(String prefix, boolean recursive) {
                 // Remove recursive result and requested dir
                 return !key.isEmpty()
                     && !Objects.equals(key, prefix)
-                    && (recursive || Path.of(key).getParent() == null);
+                    && !key.equals("/")
+                    && (recursive || Path.of(key).getParent() == null)
+                    && (includeDirectories || !key.endsWith("/"));
             });
     }