Return archive in zip stream

Signed-off-by: yichen88 <[email protected]>
powsybl · Mar 24, 2020 · 7169c9b · 7169c9b
1 parent 4ab5835
commit 7169c9b
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 21 deletions.
diff --git a/afs-core/src/main/java/com/powsybl/afs/AbstractNodeBase.java b/afs-core/src/main/java/com/powsybl/afs/AbstractNodeBase.java
@@ -9,8 +9,8 @@
 
 import com.powsybl.afs.storage.AppStorage;
 import com.powsybl.afs.storage.AppStorageArchive;
-import com.powsybl.afs.storage.Utils;
 import com.powsybl.afs.storage.NodeInfo;
+import com.powsybl.afs.storage.Utils;
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
@@ -22,6 +22,7 @@
 import java.time.ZonedDateTime;
 import java.util.*;
 import java.util.stream.Stream;
+import java.util.zip.ZipInputStream;
 
 /**
  * Base class for all node objects stored in an AFS tree.
@@ -192,6 +193,23 @@ public void archive(Path dir, boolean useZip, boolean archiveDependencies, Map<S
         }
     }
 
+    /**
+     *
+     * @param dir The directory to archive
+     * @param archiveDependencies
+     * @param outputBlackList
+     * @return Returns an ZipInputStream representing the archived file
+     */
+    public ZipInputStream openArchiveStream(Path dir, boolean archiveDependencies, Map<String, List<String>> outputBlackList) {
+        Objects.requireNonNull(dir);
+        try {
+            new AppStorageArchive(storage).archive(info.getId(), dir, archiveDependencies, outputBlackList);
+            return Utils.zip(dir, true);
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
     public void unarchive(Path dir, boolean isZipped) {
         if (isZipped) {
             String filename = dir.getFileName().toString().substring(0, dir.getFileName().toString().length() - 4);

diff --git a/afs-core/src/test/java/com/powsybl/afs/AfsBaseTest.java b/afs-core/src/test/java/com/powsybl/afs/AfsBaseTest.java
@@ -10,11 +10,7 @@
 import com.google.common.jimfs.Configuration;
 import com.google.common.jimfs.Jimfs;
 import com.powsybl.afs.mapdb.storage.MapDbAppStorage;
-import com.powsybl.afs.storage.AfsStorageException;
-import com.powsybl.afs.storage.AppStorage;
-import com.powsybl.afs.storage.InMemoryEventsBus;
-import com.powsybl.afs.storage.NodeGenericMetadata;
-import com.powsybl.afs.storage.NodeInfo;
+import com.powsybl.afs.storage.*;
 import com.powsybl.computation.ComputationManager;
 import com.powsybl.iidm.network.NetworkFactoryService;
 import org.junit.After;
@@ -31,6 +27,7 @@
 import java.nio.file.Path;
 import java.util.*;
 import java.util.function.BiConsumer;
+import java.util.zip.ZipInputStream;
 
 import static org.junit.Assert.*;
 
@@ -250,6 +247,22 @@ public void archiveAndUnarchiveTestWithZip() throws IOException {
         assertEquals(1, dir2.getChildren().size());
     }
 
+    @Test
+    public void archiveInMemory() throws IOException {
+        Project project = afs.getRootFolder().createProject("test");
+        ProjectFolder rootFolder = project.getRootFolder();
+        ProjectFolder dir1 = rootFolder.createFolder("dir1");
+        String nodeId = dir1.getId();
+        Path rootDir = fileSystem.getPath("/root");
+        Files.createDirectory(rootDir);
+        Files.createDirectory(rootDir.resolve("test"));
+        try (ZipInputStream zis = dir1.openArchiveStream(rootDir.resolve("test"), true, new HashMap<>())) {
+            assertEquals(nodeId + "/", zis.getNextEntry().getName());
+            assertEquals(nodeId + "/info.json", zis.getNextEntry().getName());
+            assertNull(zis.getNextEntry());
+        }
+    }
+
     @Test
     public void archiveAndUnarchiveTestWithDirAndBlackList() throws IOException {
         Project project = afs.getRootFolder().createProject("test");

diff --git a/afs-storage-api/src/main/java/com/powsybl/afs/storage/Utils.java b/afs-storage-api/src/main/java/com/powsybl/afs/storage/Utils.java
@@ -44,29 +44,52 @@ private Utils() throws IllegalAccessException {
      */
     public static void zip(Path dir, Path zipPath, boolean deleteDirectory) throws IOException {
         try (ZipOutputStream zos = new ZipOutputStream(Files.newOutputStream(zipPath));
-             Stream<Path> walk = Files.walk(dir);) {
-            walk.filter(someFileToZip -> !someFileToZip.equals(dir))
-                    .forEach(
-                        someFileToZip -> {
-                            Path pathInZip = dir.relativize(someFileToZip);
-                            try {
-                                if (Files.isDirectory(someFileToZip)) {
-                                    addDirectory(zos, pathInZip);
-                                } else {
-                                    addFile(zos, someFileToZip, pathInZip);
-                                }
-                            } catch (IOException e) {
-                                throw new AfsStorageException(e.getMessage());
-                            }
+             Stream<Path> walk = Files.walk(dir)) {
+            zip(walk, dir, zos);
+        } catch (IOException | AfsStorageException e) {
+            throw new IOException(e);
+        }
 
-                        });
+        if (deleteDirectory) {
+            deleteDirectory(dir);
+        }
+    }
+
+    public static ZipInputStream zip(Path dir, boolean deleteDirectory) throws IOException {
+        byte[] bytes;
+        try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
+             ZipOutputStream zos = new ZipOutputStream(baos);
+             Stream<Path> walk = Files.walk(dir)) {
+            zip(walk, dir, zos);
+            zos.close();
+            baos.close();
+            bytes = baos.toByteArray();
         } catch (IOException | AfsStorageException e) {
             throw new IOException(e);
         }
 
         if (deleteDirectory) {
             deleteDirectory(dir);
         }
+        return new ZipInputStream(new ByteArrayInputStream(bytes));
+    }
+
+    private static void zip(Stream<Path> walk, Path dir, ZipOutputStream zos) {
+        walk.filter(someFileToZip -> !someFileToZip.equals(dir))
+                .forEach(
+                    someFileToZip -> {
+                        Path pathInZip = dir.relativize(someFileToZip);
+                        try {
+                            if (Files.isDirectory(someFileToZip)) {
+                                addDirectory(zos, pathInZip);
+                            } else {
+                                addFile(zos, someFileToZip, pathInZip);
+                            }
+                        } catch (IOException e) {
+                            throw new AfsStorageException(e.getMessage());
+                        }
+
+                    });
     }
 
     /**