diff --git a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/controller/dataservice/ProjectController.java b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/controller/dataservice/ProjectController.java index ce13ef340b..b76a35fd2e 100644 --- a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/controller/dataservice/ProjectController.java +++ b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/controller/dataservice/ProjectController.java @@ -10,16 +10,12 @@ import io.swagger.v3.oas.annotations.tags.Tag; import io.swagger.v3.oas.annotations.tags.Tags; import jakarta.transaction.Transactional; -import java.io.ByteArrayOutputStream; import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; import java.util.List; import java.util.Optional; import java.util.UUID; -import java.util.zip.ZipEntry; -import java.util.zip.ZipInputStream; -import java.util.zip.ZipOutputStream; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.springframework.http.HttpHeaders; @@ -512,31 +508,13 @@ public ResponseEntity exportProject(@PathVariable("id") final UUID id) { try { final ProjectExport export = cloneService.exportProject(id); - final byte[] exportBytes = objectMapper.writeValueAsBytes(export); - - try (ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(); - ZipOutputStream zipOutputStream = new ZipOutputStream(byteArrayOutputStream)) { - - final ZipEntry zipEntry = new ZipEntry("project.json"); - zipOutputStream.putNextEntry(zipEntry); - zipOutputStream.write(exportBytes); - zipOutputStream.closeEntry(); - - zipOutputStream.finish(); - zipOutputStream.close(); - - final HttpHeaders headers = new HttpHeaders(); - headers.setContentType(MediaType.parseMediaType("application/zip")); - final String filename = "project-" + id + ".zip"; - headers.setContentDispositionFormData(filename, filename); - headers.setCacheControl("must-revalidate, post-check=0, pre-check=0"); - - return new ResponseEntity<>(byteArrayOutputStream.toByteArray(), headers, HttpStatus.OK); - } catch (final IOException e) { - e.printStackTrace(); - return new ResponseEntity<>(HttpStatus.INTERNAL_SERVER_ERROR); - } + final HttpHeaders headers = new HttpHeaders(); + headers.setContentType(MediaType.parseMediaType("application/zip")); + final String filename = "project-" + id + ".zip"; + headers.setContentDispositionFormData(filename, filename); + headers.setCacheControl("must-revalidate, post-check=0, pre-check=0"); + return new ResponseEntity<>(export.getAsZipFile(), headers, HttpStatus.OK); } catch (final Exception e) { log.error("Error exporting project", e); throw new ResponseStatusException( @@ -578,29 +556,12 @@ public ResponseEntity importProject(@RequestPart("file") final Multipar return ResponseEntity.badRequest().build(); } - final ProjectExport projectExport; + final ProjectExport projectExport = new ProjectExport(); try { - final ZipInputStream zipInputStream = new ZipInputStream(input.getInputStream()); - final ZipEntry zipEntry = zipInputStream.getNextEntry(); - final ByteArrayOutputStream outputStream = new ByteArrayOutputStream(); - - final byte[] buffer = new byte[1024]; - int length; - while ((zipEntry != null) && (length = zipInputStream.read(buffer)) > 0) { - outputStream.write(buffer, 0, length); - } - - final byte[] unzippedBytes = outputStream.toByteArray(); - - zipInputStream.closeEntry(); - zipInputStream.close(); - outputStream.close(); - - projectExport = objectMapper.readValue(unzippedBytes, ProjectExport.class); - - } catch (final IOException e) { - e.printStackTrace(); - return ResponseEntity.internalServerError().build(); + projectExport.loadFromZipFile(input.getInputStream()); + } catch (final Exception e) { + log.error("Error parsing project", e); + return ResponseEntity.badRequest().build(); } final String userId = currentUserService.get().getId(); diff --git a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/AssetExportDeserializer.java b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/AssetExportDeserializer.java index a728fadf26..47e1b80307 100644 --- a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/AssetExportDeserializer.java +++ b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/AssetExportDeserializer.java @@ -14,21 +14,21 @@ public class AssetExportDeserializer extends JsonDeserializer { @Override - public AssetExport deserialize(JsonParser jp, DeserializationContext ctxt) throws IOException { - ObjectMapper mapper = (ObjectMapper) jp.getCodec(); - JsonNode node = mapper.readTree(jp); + public AssetExport deserialize(final JsonParser jp, final DeserializationContext ctxt) throws IOException { + final ObjectMapper mapper = (ObjectMapper) jp.getCodec(); + final JsonNode node = mapper.readTree(jp); final String assetTypeStr = node.get("type").asText(); final AssetType assetType = AssetType.getAssetType(assetTypeStr, mapper); - TerariumAsset asset = mapper.treeToValue(node.get("asset"), assetType.getAssetClass()); + final TerariumAsset asset = mapper.treeToValue(node.get("asset"), assetType.getAssetClass()); Map files = new HashMap<>(); - if (node.has("files")) { + if (node.hasNonNull("files")) { files = mapper.convertValue(node.get("files"), new TypeReference>() {}); } - AssetExport export = new AssetExport(); + final AssetExport export = new AssetExport(); export.setType(assetType); export.setAsset(asset); export.setFiles(files); diff --git a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/ContentTypeDeserializer.java b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/ContentTypeDeserializer.java index 0eb91c0302..24cf7fed5a 100644 --- a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/ContentTypeDeserializer.java +++ b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/ContentTypeDeserializer.java @@ -20,7 +20,7 @@ public ContentType deserialize(final JsonParser jp, final DeserializationContext final JsonNode node = mapper.readTree(jp); final String mimeTypeStr = node.get("mimeType").asText(); - if (!node.has("charset")) { + if (!node.hasNonNull("charset")) { return ContentType.create(mimeTypeStr); } diff --git a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/FileExport.java b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/FileExport.java index 1c14a045ee..34021893ba 100644 --- a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/FileExport.java +++ b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/FileExport.java @@ -1,5 +1,6 @@ package software.uncharted.terarium.hmiserver.models.dataservice; +import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.databind.annotation.JsonDeserialize; import lombok.Data; import lombok.experimental.Accessors; @@ -12,5 +13,6 @@ public class FileExport { @JsonDeserialize(using = ContentTypeDeserializer.class) ContentType contentType; + @JsonIgnore byte[] bytes; } diff --git a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/document/DocumentAsset.java b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/document/DocumentAsset.java index d6f6cbd65e..48897a059a 100644 --- a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/document/DocumentAsset.java +++ b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/document/DocumentAsset.java @@ -72,23 +72,19 @@ public class DocumentAsset extends TerariumAsset { @Override public List getFileNames() { - final List res = new ArrayList<>(); - if (this.fileNames != null) { - for (final String fileName : fileNames) { - if (!res.contains(fileName)) { - res.add(fileName); - } - } + if (this.fileNames == null) { + this.fileNames = new ArrayList<>(); } + // ensure these are included in filenames if (this.assets != null) { for (final DocumentExtraction asset : assets) { - if (!res.contains(asset.getFileName())) { - res.add(asset.getFileName()); + if (!this.fileNames.contains(asset.getFileName())) { + this.fileNames.add(asset.getFileName()); } } } - return res; + return this.fileNames; } /** diff --git a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/project/ProjectExport.java b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/project/ProjectExport.java index a340a0f74f..a646801387 100644 --- a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/project/ProjectExport.java +++ b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/models/dataservice/project/ProjectExport.java @@ -1,5 +1,10 @@ package software.uncharted.terarium.hmiserver.models.dataservice.project; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; import java.util.ArrayList; import java.util.HashMap; import java.util.HashSet; @@ -7,10 +12,14 @@ import java.util.Map; import java.util.Set; import java.util.UUID; +import java.util.zip.ZipEntry; +import java.util.zip.ZipInputStream; +import java.util.zip.ZipOutputStream; import lombok.Data; import lombok.experimental.Accessors; import software.uncharted.terarium.hmiserver.models.TerariumAsset; import software.uncharted.terarium.hmiserver.models.dataservice.AssetExport; +import software.uncharted.terarium.hmiserver.models.dataservice.FileExport; import software.uncharted.terarium.hmiserver.utils.AssetDependencyUtil; import software.uncharted.terarium.hmiserver.utils.AssetDependencyUtil.AssetDependencyMap; @@ -21,6 +30,88 @@ public class ProjectExport { Project project; List assets = new ArrayList<>(); + private byte[] readZipEntry(final ZipInputStream zipInputStream) throws IOException { + final ByteArrayOutputStream baos = new ByteArrayOutputStream(); + final byte[] buffer = new byte[1024]; + int count; + while ((count = zipInputStream.read(buffer)) != -1) { + baos.write(buffer, 0, count); + } + return baos.toByteArray(); + } + + public void loadFromZipFile(final InputStream inputStream) throws IOException { + + final ObjectMapper objectMapper = new ObjectMapper(); + + final ZipInputStream zipInputStream = new ZipInputStream(inputStream); + + // get the project json + ZipEntry zipEntry = zipInputStream.getNextEntry(); + if (zipEntry == null || !zipEntry.getName().equals("project.json")) { + throw new IllegalArgumentException("Invalid project export file"); + } + + project = objectMapper.readValue(readZipEntry(zipInputStream), Project.class); + + // iterate on assets + while ((zipEntry = zipInputStream.getNextEntry()) != null) { + + // read the asset json + final AssetExport asset = objectMapper.readValue(readZipEntry(zipInputStream), AssetExport.class); + + // read in the file payloads + for (final Map.Entry entry : asset.getFiles().entrySet()) { + zipEntry = zipInputStream.getNextEntry(); + if (zipEntry == null) { + throw new IllegalArgumentException("Invalid project export file, expected a asset file payload"); + } + final FileExport file = entry.getValue(); + file.setBytes(readZipEntry(zipInputStream)); + } + + assets.add(asset); + } + } + + public byte[] getAsZipFile() throws JsonProcessingException, IOException { + + final ObjectMapper objectMapper = new ObjectMapper(); + + final ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(); + final ZipOutputStream zipOutputStream = new ZipOutputStream(byteArrayOutputStream); + + final byte[] projectBytes = objectMapper.writeValueAsBytes(project); + + final ZipEntry zipEntry = new ZipEntry("project.json"); + zipOutputStream.putNextEntry(zipEntry); + zipOutputStream.write(projectBytes); + zipOutputStream.closeEntry(); + + for (final AssetExport asset : assets) { + + final byte[] assetBytes = objectMapper.writeValueAsBytes(asset); + + final ZipEntry assetEntry = new ZipEntry(asset.getAsset().getId() + ".json"); + zipOutputStream.putNextEntry(assetEntry); + zipOutputStream.write(assetBytes); + zipOutputStream.closeEntry(); + + for (final Map.Entry file : asset.getFiles().entrySet()) { + final byte[] fileBytes = file.getValue().getBytes(); + + final ZipEntry fileEntry = new ZipEntry(asset.getAsset().getId() + "/" + file.getKey()); + zipOutputStream.putNextEntry(fileEntry); + zipOutputStream.write(fileBytes); + zipOutputStream.closeEntry(); + } + } + zipOutputStream.finish(); + zipOutputStream.close(); + + return byteArrayOutputStream.toByteArray(); + } + public ProjectExport clone() { final ProjectExport cloned = new ProjectExport(); diff --git a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/repository/data/ProjectRepository.java b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/repository/data/ProjectRepository.java index 2d8290a712..4821011756 100644 --- a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/repository/data/ProjectRepository.java +++ b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/repository/data/ProjectRepository.java @@ -50,7 +50,7 @@ public interface ProjectRepository extends PSCrudRepository, JpaS + "where " + " p.id in (:ids) " + " and p.deletedOn is null") - List findByIdsWithAssets(final List ids); + List findByIdsWithAssets(@Param("ids") final List ids); @Query(value = "SELECT public_asset FROM project WHERE id = :id", nativeQuery = true) Optional findPublicAssetByIdNative(@Param("id") UUID id); diff --git a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/service/data/ModelService.java b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/service/data/ModelService.java index e621812f65..ea671fc37d 100644 --- a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/service/data/ModelService.java +++ b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/service/data/ModelService.java @@ -142,11 +142,10 @@ public Model createAsset(final Model asset, final UUID projectId, final Schema.P // Execute the update request uploadEmbeddings(created.getId(), embeddings, hasWritePermission); } catch (final Exception e) { - log.error("Failed to update embeddings for document {}", created.getId(), e); + log.error("Failed to update embeddings for model {}", created.getId(), e); } }) .start(); - ; } return created; @@ -181,11 +180,10 @@ public Optional updateAsset( // Execute the update request uploadEmbeddings(updated.getId(), embeddings, hasWritePermission); } catch (final Exception e) { - log.error("Failed to update embeddings for document {}", updated.getId(), e); + log.error("Failed to update embeddings for model {}", updated.getId(), e); } }) .start(); - ; } return updatedOptional; diff --git a/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/DocumentServiceTests.java b/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/DocumentServiceTests.java index 42896e20fb..41276c1b1f 100644 --- a/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/DocumentServiceTests.java +++ b/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/DocumentServiceTests.java @@ -96,7 +96,7 @@ public void testItCanCreateDocument() throws Exception { Assertions.assertEquals(before.getId(), after.getId()); Assertions.assertNotNull(after.getId()); Assertions.assertNotNull(after.getCreatedOn()); - Assertions.assertEquals(after.getFileNames().size(), 2); + Assertions.assertEquals(3, after.getFileNames().size()); Assertions.assertNotNull(after.getGrounding()); Assertions.assertNotNull(after.getGrounding().getId());