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

part of kestra-io/kestra#2368
kestra-io · Dec 5, 2023 · 8c57d10 · 8c57d10
1 parent 2e283ea
commit 8c57d10
Showing 1 changed file with 31 additions and 8 deletions.
diff --git a/src/main/java/io/kestra/storage/gcs/GcsStorage.java b/src/main/java/io/kestra/storage/gcs/GcsStorage.java
@@ -15,8 +15,10 @@
 import java.nio.ByteBuffer;
 import java.nio.channels.Channels;
 import java.nio.channels.ReadableByteChannel;
+import java.nio.file.Path;
 import java.util.*;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import jakarta.inject.Inject;
 import jakarta.inject.Singleton;
@@ -86,18 +88,23 @@ public InputStream get(String tenantId, URI uri) throws IOException {
         }
     }
 
+    @Override
+    public List<URI> filesByPrefix(String tenantId, URI prefix) {
+        String path = getPath(tenantId, prefix);
+        return blobsForPrefix(path, true)
+            .map(BlobInfo::getName)
+            // keep only files
+            .filter(blobPath -> !blobPath.endsWith("/"))
+            .map(blobPath -> URI.create("kestra://" + prefix.getPath() + blobPath.substring(path.length())))
+            .toList();
+    }
+
     @Override
     public List<FileAttributes> list(String tenantId, URI uri) throws IOException {
         String path = getPath(tenantId, uri);
         String prefix = (path.endsWith("/")) ? path : path + "/";
-        Page<Blob> blobs = this.storage.list(config.bucket, Storage.BlobListOption.prefix(prefix),
-            Storage.BlobListOption.currentDirectory());
-        List<FileAttributes> list = blobs.streamAll()
-            .filter(blob -> {
-                String key = blob.getName().substring(prefix.length());
-                // Remove recursive result and requested dir
-                return !key.isEmpty() && !Objects.equals(key, prefix) && new File(key).getParent() == null;
-            })
+
+        List<FileAttributes> list = blobsForPrefix(prefix, false)
             .map(throwFunction(this::getGcsFileAttributes))
             .toList();
         if(list.isEmpty()) {
@@ -107,6 +114,22 @@ public List<FileAttributes> list(String tenantId, URI uri) throws IOException {
         return list;
     }
 
+    private Stream<Blob> blobsForPrefix(String prefix, boolean recursive) {
+        Storage.BlobListOption[] blobListOptions = Stream.concat(
+            Stream.of(Storage.BlobListOption.prefix(prefix)),
+            recursive ? Stream.empty() : Stream.of(Storage.BlobListOption.currentDirectory())
+        ).toArray(Storage.BlobListOption[]::new);
+        Page<Blob> blobs = this.storage.list(config.bucket, blobListOptions);
+        return blobs.streamAll()
+            .filter(blob -> {
+                String key = blob.getName().substring(prefix.length());
+                // Remove recursive result and requested dir
+                return !key.isEmpty()
+                    && !Objects.equals(key, prefix)
+                    && (recursive || Path.of(key).getParent() == null);
+            });
+    }
+
     @Override
     public boolean exists(String tenantId, URI uri) {
         try {