fix(core): implement filePathsByPrefix for search by file name in editor

part of kestra-io/kestra#2368
kestra-io · Dec 6, 2023 · b9607b9 · b9607b9
1 parent 2c21b52
commit b9607b9
Showing 1 changed file with 30 additions and 13 deletions.
diff --git a/src/main/java/io/kestra/storage/s3/S3Storage.java b/src/main/java/io/kestra/storage/s3/S3Storage.java
@@ -29,6 +29,7 @@
 import java.util.Optional;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executors;
+import java.util.stream.Stream;
 
 import static io.kestra.core.utils.Rethrow.throwFunction;
 
@@ -64,6 +65,16 @@ public InputStream get(String tenantId, URI uri) throws IOException {
         return get(path);
     }
 
+    @Override
+    public List<URI> filesByPrefix(String tenantId, URI prefix) {
+        String path = getPath(tenantId, prefix);
+        return keysForPrefix(path, true)
+            // keep only files
+            .filter(key -> !key.endsWith("/"))
+            .map(key -> URI.create("kestra://" + prefix.getPath() + key.substring(path.length())))
+            .toList();
+    }
+
     private InputStream get(String path) throws IOException {
         try {
             GetObjectRequest request = GetObjectRequest.builder()
@@ -91,21 +102,10 @@ public List<FileAttributes> list(String tenantId, URI uri) throws IOException {
         String path = getPath(tenantId, uri);
         String prefix = path.endsWith("/") ? path : path + "/";
         try {
-            ListObjectsV2Request request = ListObjectsV2Request.builder()
-                .bucket(s3Config.getBucket())
-                .prefix(prefix)
-                .build();
-            List<S3Object> contents = s3Client.listObjectsV2(request).contents();
-            List<FileAttributes> list = contents.stream()
-                .map(S3Object::key)
-                .filter(key -> {
-                    key = key.substring(prefix.length());
-                    // Remove recursive result and requested dir
-                    return !key.isEmpty() && !Objects.equals(key, prefix) && Path.of(key).getParent() == null;
-                })
+            List<FileAttributes> list = keysForPrefix(prefix, false)
                 .map(throwFunction(this::getFileAttributes))
                 .toList();
-            if(list.isEmpty()) {
+            if (list.isEmpty()) {
                 // this will throw FileNotFound if there is no directory
                 this.getAttributes(tenantId, uri);
             }
@@ -117,6 +117,23 @@ public List<FileAttributes> list(String tenantId, URI uri) throws IOException {
         }
     }
 
+    private Stream<String> keysForPrefix(String prefix, boolean recursive) {
+        ListObjectsV2Request request = ListObjectsV2Request.builder()
+            .bucket(s3Config.getBucket())
+            .prefix(prefix)
+            .build();
+        List<S3Object> contents = s3Client.listObjectsV2(request).contents();
+        return contents.stream()
+            .map(S3Object::key)
+            .filter(key -> {
+                key = key.substring(prefix.length());
+                // Remove recursive result and requested dir
+                return !key.isEmpty()
+                    && !Objects.equals(key, prefix)
+                    && (recursive || Path.of(key).getParent() == null);
+            });
+    }
+
     @Override
     public Long size(String tenantId, URI uri) throws IOException {
         try {