From 2ce0b4e07f9d728bf77e18f223bb92309b3f0a0d Mon Sep 17 00:00:00 2001 From: Ian Robinson Date: Fri, 18 Sep 2020 23:35:36 +0100 Subject: [PATCH] Neptune export updates (#81) * Add ability to replace or append S3 key * Allow user to suplpy export ID --- .../cli/PropertyGraphTargetModule.java | 7 ++++- .../services/neptune/cli/RdfTargetModule.java | 7 ++++- .../neptune/export/NeptuneExportService.java | 13 +++++++-- .../services/neptune/io/Directories.java | 9 +++---- .../services/neptune/util/S3ObjectInfo.java | 13 +++++++++ .../neptune/util/S3ObjectInfoTest.java | 27 +++++++++++++++++++ 6 files changed, 66 insertions(+), 10 deletions(-) diff --git a/neptune-export/src/main/java/com/amazonaws/services/neptune/cli/PropertyGraphTargetModule.java b/neptune-export/src/main/java/com/amazonaws/services/neptune/cli/PropertyGraphTargetModule.java index cfee769a..9fe7a1a3 100644 --- a/neptune-export/src/main/java/com/amazonaws/services/neptune/cli/PropertyGraphTargetModule.java +++ b/neptune-export/src/main/java/com/amazonaws/services/neptune/cli/PropertyGraphTargetModule.java @@ -23,6 +23,7 @@ import java.io.File; import java.io.IOException; +import java.util.UUID; public class PropertyGraphTargetModule implements RequiresMetadata, CommandWriter { @@ -54,8 +55,12 @@ public class PropertyGraphTargetModule implements RequiresMetadata, CommandWrite @Once private String region; + @Option(name = {"--export-id"}, description = "Export ID", hidden = true) + @Once + private String exportId = UUID.randomUUID().toString().replace("-", ""); + public Directories createDirectories(DirectoryStructure directoryStructure) throws IOException { - return Directories.createFor(directoryStructure, directory, tag ); + return Directories.createFor(directoryStructure, directory, exportId, tag ); } public PropertyGraphTargetConfig config(Directories directories, boolean includeTypeDefinitions){ diff --git a/neptune-export/src/main/java/com/amazonaws/services/neptune/cli/RdfTargetModule.java b/neptune-export/src/main/java/com/amazonaws/services/neptune/cli/RdfTargetModule.java index 8d873ca6..30a918f8 100644 --- a/neptune-export/src/main/java/com/amazonaws/services/neptune/cli/RdfTargetModule.java +++ b/neptune-export/src/main/java/com/amazonaws/services/neptune/cli/RdfTargetModule.java @@ -23,6 +23,7 @@ import java.io.File; import java.io.IOException; +import java.util.UUID; public class RdfTargetModule implements CommandWriter { @@ -54,8 +55,12 @@ public class RdfTargetModule implements CommandWriter { @Once private String region; + @Option(name = {"--export-id"}, description = "Export ID", hidden = true) + @Once + private String exportId = UUID.randomUUID().toString().replace("-", ""); + public Directories createDirectories(DirectoryStructure directoryStructure) throws IOException { - return Directories.createFor(directoryStructure, directory, tag ); + return Directories.createFor(directoryStructure, directory, exportId, tag ); } public RdfTargetConfig config(Directories directories) { diff --git a/neptune-export/src/main/java/com/amazonaws/services/neptune/export/NeptuneExportService.java b/neptune-export/src/main/java/com/amazonaws/services/neptune/export/NeptuneExportService.java index 995f4220..c61790e9 100644 --- a/neptune-export/src/main/java/com/amazonaws/services/neptune/export/NeptuneExportService.java +++ b/neptune-export/src/main/java/com/amazonaws/services/neptune/export/NeptuneExportService.java @@ -21,6 +21,7 @@ import com.fasterxml.jackson.databind.node.JsonNodeFactory; import com.fasterxml.jackson.databind.node.ObjectNode; import org.apache.commons.io.FileUtils; +import org.apache.commons.io.FilenameUtils; import org.apache.commons.lang.StringUtils; import java.io.*; @@ -232,9 +233,17 @@ private void uploadCompletionFileToS3(TransferManager transferManager, throw new RuntimeException("Error while writing completion file payload", e); } - S3ObjectInfo completionFileS3ObjectInfo = new S3ObjectInfo(completionFileS3Path).withNewKeySuffix(completionFile.getName()); + S3ObjectInfo completionFileS3ObjectInfo = + new S3ObjectInfo(completionFileS3Path).replaceOrAppendKey( + "_COMPLETION_ID_", + FilenameUtils.getBaseName(completionFile.getName()), + completionFile.getName()); + + Upload upload = transferManager.upload( + completionFileS3ObjectInfo.bucket(), + completionFileS3ObjectInfo.key(), + completionFile); - Upload upload = transferManager.upload(completionFileS3ObjectInfo.bucket(), completionFileS3ObjectInfo.key(), completionFile); try { upload.waitForUploadResult(); } catch (InterruptedException e) { diff --git a/neptune-export/src/main/java/com/amazonaws/services/neptune/io/Directories.java b/neptune-export/src/main/java/com/amazonaws/services/neptune/io/Directories.java index 8ca9d6e5..0551ae8b 100644 --- a/neptune-export/src/main/java/com/amazonaws/services/neptune/io/Directories.java +++ b/neptune-export/src/main/java/com/amazonaws/services/neptune/io/Directories.java @@ -18,7 +18,6 @@ import java.io.File; import java.io.IOException; -import java.net.URI; import java.nio.file.Files; import java.nio.file.Path; import java.util.Collection; @@ -29,16 +28,14 @@ public class Directories { private static final String CONFIG_FILE = "config.json"; private static final String QUERIES_FILE = "queries.json"; - public static Directories createFor(DirectoryStructure directoryStructure, File root, String tag) throws IOException { + public static Directories createFor(DirectoryStructure directoryStructure, File root, String exportId, String tag) throws IOException { if (root == null) { throw new IllegalArgumentException("You must supply a directory"); } - String uuid = UUID.randomUUID().toString().replace("-", ""); - String directoryName = tag.isEmpty() ? - uuid : - String.format("%s-%s", tag, uuid); + exportId : + String.format("%s-%s", tag, exportId); Path rootDirectory = root.toPath(); Path directory = rootDirectory.resolve(directoryName); diff --git a/neptune-export/src/main/java/com/amazonaws/services/neptune/util/S3ObjectInfo.java b/neptune-export/src/main/java/com/amazonaws/services/neptune/util/S3ObjectInfo.java index faf21fa6..95609f32 100644 --- a/neptune-export/src/main/java/com/amazonaws/services/neptune/util/S3ObjectInfo.java +++ b/neptune-export/src/main/java/com/amazonaws/services/neptune/util/S3ObjectInfo.java @@ -44,6 +44,19 @@ public S3ObjectInfo withNewKeySuffix(String suffix) { return new S3ObjectInfo( String.format("s3://%s/%s", bucket, new File(key, suffix).getPath())); } + public S3ObjectInfo replaceOrAppendKey(String placeholder, String ifPresent, String ifAbsent) { + + File file = key.contains(placeholder) ? + new File(key.replace(placeholder, ifPresent)) : + new File(key, ifAbsent); + + return new S3ObjectInfo( String.format("s3://%s/%s", bucket, file.getPath())); + } + + public S3ObjectInfo replaceOrAppendKey(String placeholder, String ifPresent) { + return replaceOrAppendKey(placeholder, ifPresent, ifPresent); + } + @Override public String toString() { return String.format("s3://%s/%s", bucket, key); diff --git a/neptune-export/src/test/java/com/amazonaws/services/neptune/util/S3ObjectInfoTest.java b/neptune-export/src/test/java/com/amazonaws/services/neptune/util/S3ObjectInfoTest.java index 861d839b..28b2be1d 100644 --- a/neptune-export/src/test/java/com/amazonaws/services/neptune/util/S3ObjectInfoTest.java +++ b/neptune-export/src/test/java/com/amazonaws/services/neptune/util/S3ObjectInfoTest.java @@ -93,4 +93,31 @@ public void canCreateNewKeyForKeyWithTrailingSlash() { assertEquals("a/b/c/dir", s3ObjectInfo.withNewKeySuffix("dir").key()); } + + @Test + public void canReplacePlaceholderInKey() { + String s3Uri = "s3://my-bucket/a/b/_COMPLETION_ID_/manifest.json"; + + S3ObjectInfo s3ObjectInfo = new S3ObjectInfo(s3Uri); + + assertEquals("a/b/123/manifest.json", s3ObjectInfo.replaceOrAppendKey("_COMPLETION_ID_", "123").key()); + } + + @Test + public void canAppendSuffixIfNoPlaceholder() { + String s3Uri = "s3://my-bucket/a/b/"; + + S3ObjectInfo s3ObjectInfo = new S3ObjectInfo(s3Uri); + + assertEquals("a/b/123", s3ObjectInfo.replaceOrAppendKey("_COMPLETION_ID_", "123").key()); + } + + @Test + public void canAppendAltSuffixIfNoPlaceholder() { + String s3Uri = "s3://my-bucket/a/b/"; + + S3ObjectInfo s3ObjectInfo = new S3ObjectInfo(s3Uri); + + assertEquals("a/b/123.json", s3ObjectInfo.replaceOrAppendKey("_COMPLETION_ID_", "123", "123.json").key()); + } } \ No newline at end of file