Modify attachment service API to decouple from job id

Previously required to use a pre-generated job id as the key. In some ways this made sense but it forced generation of a job id in the controller tier when it makes more sense to do that in the service tier closer to the source of truth (database). This change makes it so attachments can be saved and a unique identifier is returned to the caller to be used as a token to get the attachments.

genie-web/src/main/java/com/netflix/genie/web/services/AttachmentService.java

@@ -21,10 +21,13 @@
 import org.springframework.validation.annotation.Validated;
 
 import java.io.File;
+import java.io.IOException;
 import java.io.InputStream;
+import java.nio.file.Path;
+import java.util.Map;
 
 /**
- * APIs for dealing with attachments sent in with Genie jobs. Implementations will handle where to store them and
+ * APIs for dealing with attachments sent in with Genie requests. Implementations will handle where to store them and
  * how to retrieve them when requested.
  *
  * @author tgianos
@@ -40,7 +43,9 @@ public interface AttachmentService {
      * @param filename The name of the attachment
      * @param content  A stream to access the contents of the attachment
      * @throws GenieException For any error during the save process
+     * @deprecated Use {@link #saveAll(Map)} instead
      */
+    @Deprecated
     void save(String jobId, String filename, InputStream content) throws GenieException;
 
     /**
@@ -49,14 +54,46 @@ public interface AttachmentService {
      * @param jobId       The id of the job to get the attachments for.
      * @param destination The directory to copy the attachments into
      * @throws GenieException For any error during the copy process
+     * @deprecated Use {@link #copyAll(String, Path)} instead
      */
+    @Deprecated
     void copy(String jobId, File destination) throws GenieException;
 
     /**
      * Delete the attachments for the given job.
      *
      * @param jobId The id of the job to delete the attachments for
      * @throws GenieException For any error during the delete process
+     * @deprecated Use {@link #deleteAll(String)} instead
      */
+    @Deprecated
     void delete(String jobId) throws GenieException;
+
+    /**
+     * Save all attachments for a given request.
+     *
+     * @param attachments The map of filename to contents for all attachments. All input streams will be closed after
+     *                    this method returns
+     * @return A unique identifier that can be used to reference the attachments later
+     * @throws IOException If unable to save any of the attachments
+     */
+    String saveAll(Map<String, InputStream> attachments) throws IOException;
+
+    /**
+     * Copy all attachments associated with the given id into the provided {@literal destination}.
+     *
+     * @param id          The id that was returned from the original call to {@link #saveAll(Map)}
+     * @param destination The destination where the attachments should be copied. Must be a directory if it already
+     *                    exists. If it doesn't exist it will be created.
+     * @throws IOException If the copy fails for any reason
+     */
+    void copyAll(String id, Path destination) throws IOException;
+
+    /**
+     * Delete all the attachments that were associated with the given {@literal id}.
+     *
+     * @param id The id that was returned from the original call to {@link #saveAll(Map)}
+     * @throws IOException On error during deletion
+     */
+    void deleteAll(String id) throws IOException;
 }

genie-web/src/main/java/com/netflix/genie/web/services/impl/FileSystemAttachmentService.java

@@ -17,6 +17,7 @@
  */
 package com.netflix.genie.web.services.impl;
 
+import com.google.common.collect.ImmutableMap;
 import com.netflix.genie.common.exceptions.GenieException;
 import com.netflix.genie.common.exceptions.GeniePreconditionException;
 import com.netflix.genie.common.exceptions.GenieServerException;
@@ -30,6 +31,10 @@
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.Map;
+import java.util.UUID;
 
 /**
  * Implementation of the AttachmentService interface which saves and retrieves attachments from the local filesystem.
@@ -40,15 +45,24 @@
 @Slf4j
 public class FileSystemAttachmentService implements AttachmentService {
 
-    private File attachmentDirectory;
+    private final Path attachmentDirectory;
 
     /**
      * Constructor.
      *
      * @param attachmentsDirectory The directory to use or null if want to default to system temp directory
      */
     public FileSystemAttachmentService(final String attachmentsDirectory) {
-        this.createAttachmentDirectory(attachmentsDirectory);
+        String attachmentsDirectoryPath = attachmentsDirectory;
+        if (!attachmentsDirectoryPath.endsWith(File.separator)) {
+            attachmentsDirectoryPath = attachmentsDirectory + File.separator;
+        }
+
+        try {
+            this.attachmentDirectory = Files.createDirectories(Paths.get(new URI(attachmentsDirectoryPath)));
+        } catch (final IOException | URISyntaxException e) {
+            throw new IllegalArgumentException("Unable to create attachment directory: " + attachmentsDirectoryPath, e);
+        }
     }
 
     /**
@@ -60,10 +74,8 @@ public void save(
         final String filename,
         final InputStream content
     ) throws GenieException {
-        final File attachment = new File(attachmentDirectory, jobId + "/" + filename);
         try {
-            FileUtils.copyInputStreamToFile(content, attachment);
-            log.info("Saved {} to {}", filename, attachment.getAbsolutePath());
+            this.writeAttachments(jobId, ImmutableMap.of(filename, content));
         } catch (final IOException ioe) {
             throw new GenieServerException("Failed to save attachment", ioe);
         }
@@ -77,13 +89,10 @@ public void copy(final String jobId, final File destination) throws GenieExcepti
         if (destination.exists() && !destination.isDirectory()) {
             throw new GeniePreconditionException(destination + " is not a directory and it needs to be.");
         }
-        final File source = new File(attachmentDirectory, jobId);
-        if (source.exists() && source.isDirectory()) {
-            try {
-                FileUtils.copyDirectory(source, destination);
-            } catch (final IOException ioe) {
-                throw new GenieServerException("Failed to copy attachment directory", ioe);
-            }
+        try {
+            this.copyAll(jobId, destination.toPath());
+        } catch (final IOException ioe) {
+            throw new GenieServerException("Failed to copy attachment directory", ioe);
         }
     }
 
@@ -92,32 +101,56 @@ public void copy(final String jobId, final File destination) throws GenieExcepti
      */
     @Override
     public void delete(final String jobId) throws GenieException {
-        final File jobDir = new File(attachmentDirectory, jobId);
-        if (jobDir.exists()) {
-            try {
-                FileUtils.deleteDirectory(jobDir);
-            } catch (final IOException ioe) {
-                throw new GenieServerException("Failed to delete directory " + jobId, ioe);
-            }
+        try {
+            this.deleteAll(jobId);
+        } catch (final IOException ioe) {
+            throw new GenieServerException("Failed to delete directory " + jobId, ioe);
         }
     }
 
-    private void createAttachmentDirectory(final String attachmentsDirectory) {
-        String attachmentsDirectoryPath = attachmentsDirectory;
-        if (!attachmentsDirectoryPath.endsWith(File.separator)) {
-            attachmentsDirectoryPath = attachmentsDirectory + File.separator;
+    /**
+     * {@inheritDoc}
+     */
+    @Override
+    public String saveAll(final Map<String, InputStream> attachments) throws IOException {
+        final String requestId = UUID.randomUUID().toString();
+        this.writeAttachments(requestId, attachments);
+        return requestId;
+    }
+
+    /**
+     * {@inheritDoc}
+     */
+    @Override
+    public void copyAll(final String id, final Path destination) throws IOException {
+        Files.createDirectories(destination);
+        final Path attachmentDir = this.attachmentDirectory.resolve(id);
+        if (Files.exists(attachmentDir) && Files.isDirectory(attachmentDir)) {
+            // Lets use this cause I don't feel like writing a visitor
+            FileUtils.copyDirectory(attachmentDir.toFile(), destination.toFile());
         }
-        try {
-            final File dir = new File(new URI(attachmentsDirectoryPath));
-            if (!dir.exists()) {
-                Files.createDirectories(dir.toPath());
+    }
+
+    /**
+     * {@inheritDoc}
+     */
+    @Override
+    public void deleteAll(final String id) throws IOException {
+        final Path attachmentDir = this.attachmentDirectory.resolve(id);
+        FileUtils.deleteDirectory(attachmentDir.toFile());
+    }
+
+    private void writeAttachments(final String id, final Map<String, InputStream> attachments) throws IOException {
+        final Path requestDir = Files.createDirectories(this.attachmentDirectory.resolve(id));
+
+        for (final Map.Entry<String, InputStream> entry : attachments.entrySet()) {
+            // Sanitize the filename
+            final Path fileName = Paths.get(entry.getKey()).getFileName();
+            final Path file = requestDir.resolve(fileName);
+            try (InputStream contents = entry.getValue()) {
+                final long byteCount = Files.copy(contents, file);
+                log.debug("Wrote {} bytes for attachment {} to {}", byteCount, fileName, file);
             }
-            this.attachmentDirectory = dir;
-        } catch (IOException | URISyntaxException e) {
-            throw new IllegalArgumentException(
-                "Failed to create attachments directory " + attachmentsDirectoryPath,
-                e
-            );
         }
     }
 }

genie-web/src/test/java/com/netflix/genie/web/services/impl/FileSystemAttachmentServiceTest.java

@@ -17,10 +17,12 @@
  */
 package com.netflix.genie.web.services.impl;
 
+import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
 import com.netflix.genie.common.exceptions.GenieException;
 import com.netflix.genie.common.exceptions.GeniePreconditionException;
 import org.apache.commons.io.FileUtils;
+import org.hamcrest.Matchers;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Rule;
@@ -31,7 +33,12 @@
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
+import java.io.InputStream;
 import java.nio.charset.Charset;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Map;
 import java.util.Set;
 import java.util.UUID;
 
@@ -80,10 +87,9 @@ public void canSaveAttachment() throws GenieException, IOException {
      * Make sure it can't copy if the destination isn't a directory.
      *
      * @throws GenieException on error
-     * @throws IOException    on error
      */
     @Test(expected = GeniePreconditionException.class)
-    public void cantCopyIfDestinationIsntDirectory() throws GenieException, IOException {
+    public void cantCopyIfDestinationIsntDirectory() throws GenieException {
         final File destination = Mockito.mock(File.class);
         Mockito.when(destination.exists()).thenReturn(true);
         Mockito.when(destination.isDirectory()).thenReturn(false);
@@ -132,6 +138,55 @@ public void canDeleteAttachments() throws GenieException, IOException {
         attachments.forEach(file -> Assert.assertFalse(file.exists()));
     }
 
+    /**
+     * Test that all exceptions can be saved at once.
+     *
+     * @throws IOException on error
+     */
+    @Test
+    public void testAttachmentLifecycle() throws IOException {
+        final Path sourceDirectory = this.folder.newFolder().toPath();
+        final int numAttachments = 10;
+        final Set<Path> sourceAttachments = this.createAttachments(sourceDirectory, numAttachments);
+        final Map<String, InputStream> attachments = Maps.newHashMap();
+        for (final Path attachment : sourceAttachments) {
+            final Path filename = attachment.getFileName();
+            if (filename != null) {
+                attachments.put(filename.toString(), Files.newInputStream(attachment));
+            }
+        }
+        final String id = this.service.saveAll(attachments);
+        final Path expectedAttachmentDirectory = this.folder.getRoot().toPath().resolve(id);
+        Assert.assertTrue(Files.exists(expectedAttachmentDirectory));
+        Assert.assertThat(Files.list(expectedAttachmentDirectory).count(), Matchers.is((long) numAttachments));
+        for (final Path attachment : sourceAttachments) {
+            Assert.assertTrue(
+                FileUtils.contentEquals(
+                    attachment.toFile(),
+                    expectedAttachmentDirectory.resolve(attachment.getFileName()).toFile()
+                )
+            );
+        }
+
+        // Attachments saved successfully lets try to copy
+        final Path copyDirectory = this.folder.newFolder().toPath();
+        this.service.copyAll(id, copyDirectory);
+        Assert.assertTrue(Files.exists(copyDirectory));
+        Assert.assertThat(Files.list(copyDirectory).count(), Matchers.is((long) numAttachments));
+        for (final Path attachment : sourceAttachments) {
+            Assert.assertTrue(
+                FileUtils.contentEquals(
+                    attachment.toFile(),
+                    copyDirectory.resolve(attachment.getFileName()).toFile()
+                )
+            );
+        }
+
+        // Delete Them
+        this.service.deleteAll(id);
+        Assert.assertFalse(Files.exists(expectedAttachmentDirectory));
+    }
+
     private Set<File> saveAttachments(final String jobId) throws GenieException, IOException {
         final Set<File> attachments = Sets.newHashSet();
         for (int i = 0; i < 10; i++) {
@@ -163,4 +218,15 @@ private File saveAttachment(final String jobId) throws GenieException, IOExcepti
             }
         }
     }
+
+    private Set<Path> createAttachments(final Path targetDirectory, final int numAttachments) throws IOException {
+        final Set<Path> attachments = Sets.newHashSet();
+        for (int i = 0; i < numAttachments; i++) {
+            final Path attachment = targetDirectory.resolve(UUID.randomUUID().toString());
+            Files.write(attachment, ("Select * FROM my_table where id = " + i + ";").getBytes(StandardCharsets.UTF_8));
+            attachments.add(attachment);
+        }
+
+        return attachments;
+    }
 }