Neptune export updates (#81)

* Add ability to replace or append S3 key * Allow user to suplpy export ID
awslabs · Sep 18, 2020 · 2ce0b4e · 2ce0b4e
1 parent 6d347da
commit 2ce0b4e
Show file tree

Hide file tree

Showing 6 changed files with 66 additions and 10 deletions.
diff --git a/...ne-export/src/main/java/com/amazonaws/services/neptune/cli/PropertyGraphTargetModule.java b/...ne-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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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());
+    }
 }