feat(core): ScriptService now allow consuming (internalStorage, local…

…File) when replacing in commands
kestra-io · Mar 25, 2024 · ec06079 · ec06079
1 parent 6f3039e
commit ec06079
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 9 deletions.
diff --git a/core/src/main/java/io/kestra/core/models/script/ScriptService.java b/core/src/main/java/io/kestra/core/models/script/ScriptService.java
@@ -15,6 +15,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.BiConsumer;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -32,19 +33,33 @@ private ScriptService() {
     }
 
     public static String replaceInternalStorage(RunContext runContext, @Nullable String command) throws IOException {
+        return ScriptService.replaceInternalStorage(runContext, command, (s, s2) -> {});
+    }
+
+    public static String replaceInternalStorage(RunContext runContext, @Nullable String command, BiConsumer<String, String> internalStorageToLocalFileConsumer) throws IOException {
         if (command == null) {
             return "";
         }
 
         return INTERNAL_STORAGE_PATTERN
             .matcher(command)
-            .replaceAll(throwFunction(matchResult -> saveOnLocalStorage(runContext, matchResult.group())));
+            .replaceAll(throwFunction(matchResult -> {
+                String localFile = saveOnLocalStorage(runContext, matchResult.group());
+
+                internalStorageToLocalFileConsumer.accept(matchResult.group(), localFile);
+
+                return localFile;
+            }));
     }
 
     public static List<String> uploadInputFiles(RunContext runContext, List<String> commands) throws IOException {
+        return ScriptService.uploadInputFiles(runContext, commands, (s, s2) -> {});
+    }
+
+    public static List<String> uploadInputFiles(RunContext runContext, List<String> commands, BiConsumer<String, String> internalStorageToLocalFileConsumer) throws IOException {
         return commands
             .stream()
-            .map(throwFunction(s -> replaceInternalStorage(runContext, s)))
+            .map(throwFunction(s -> replaceInternalStorage(runContext, s, internalStorageToLocalFileConsumer)))
             .collect(Collectors.toList());
 
     }

diff --git a/core/src/test/java/io/kestra/core/models/script/ScriptServiceTest.java b/core/src/test/java/io/kestra/core/models/script/ScriptServiceTest.java
@@ -9,7 +9,10 @@
 import java.net.URI;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.*;
@@ -31,23 +34,39 @@ void replaceInternalStorage() throws IOException {
         if (!path.toFile().exists()) {
             Files.createFile(path);
         }
-        command = ScriptService.replaceInternalStorage(runContext, "my command with a file: kestra://some/file.txt");
-        assertThat(command, startsWith("my command with a file: /tmp/"));
-        path.toFile().delete();
+
+        String internalStorageUri = "kestra://some/file.txt";
+        AtomicReference<String> localFile = new AtomicReference<>();
+        try {
+            command = ScriptService.replaceInternalStorage(runContext, "my command with a file: " + internalStorageUri, (ignored, file) -> localFile.set(file));
+            assertThat(command, is("my command with a file: " + localFile.get()));
+            assertThat(Path.of(localFile.get()).toFile().exists(), is(true));
+        } finally {
+            Path.of(localFile.get()).toFile().delete();
+            path.toFile().delete();
+        }
     }
 
     @Test
     void uploadInputFiles() throws IOException {
         var runContext = runContextFactory.of();
+
         Path path = Path.of("/tmp/unittest/file.txt");
         if (!path.toFile().exists()) {
             Files.createFile(path);
         }
 
-        var commands  = ScriptService.uploadInputFiles(runContext, List.of("my command with a file: kestra://some/file.txt"));
-        assertThat(commands, not(empty()));
-        assertThat(commands.get(0), startsWith("my command with a file: /tmp/"));
-        path.toFile().delete();
+        Map<String, String> localFileByInternalStorage = new HashMap<>();
+        String internalStorageUri = "kestra://some/file.txt";
+        try {
+            var commands = ScriptService.uploadInputFiles(runContext, List.of("my command with a file: " + internalStorageUri), localFileByInternalStorage::put);
+            assertThat(commands, not(empty()));
+            assertThat(commands.get(0), is("my command with a file: " + localFileByInternalStorage.get(internalStorageUri)));
+            assertThat(Path.of(localFileByInternalStorage.get(internalStorageUri)).toFile().exists(), is(true));
+        } finally {
+            localFileByInternalStorage.forEach((k, v) -> Path.of(v).toFile().delete());
+            path.toFile().delete();
+        }
     }
 
     @Test