NIFI-8200: Modifying PutAzureDataLakeStorage to delete temp file if e…

…xception was thrown in uploadContent()
apache · Feb 10, 2021 · 81c2dec · 81c2dec
1 parent d5d5207
commit 81c2dec
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/...ssors/src/main/java/org/apache/nifi/processors/azure/storage/PutAzureDataLakeStorage.java b/...ssors/src/main/java/org/apache/nifi/processors/azure/storage/PutAzureDataLakeStorage.java
@@ -126,6 +126,9 @@ public void onTrigger(final ProcessContext context, final ProcessSession session
                 if (length > 0) {
                     try (final InputStream rawIn = session.read(flowFile); final BufferedInputStream bufferedIn = new BufferedInputStream(rawIn)) {
                         uploadContent(fileClient, bufferedIn, length);
+                    } catch (Exception e) {
+                        fileClient.delete();
+                        throw e;
                     }
                 }
 

diff --git a/...ors/src/test/java/org/apache/nifi/processors/azure/storage/ITPutAzureDataLakeStorage.java b/...ors/src/test/java/org/apache/nifi/processors/azure/storage/ITPutAzureDataLakeStorage.java
@@ -29,6 +29,7 @@
 import org.junit.Ignore;
 import org.junit.Test;
 
+import java.io.InputStream;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
@@ -43,6 +44,11 @@
 import static org.apache.nifi.processors.azure.storage.utils.ADLSAttributes.ATTR_NAME_PRIMARY_URI;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 
 public class ITPutAzureDataLakeStorage extends AbstractAzureDataLakeStorageIT {
 
@@ -241,6 +247,17 @@ public void testPutFileWithELButFileNameIsNotSpecified() {
         assertFailure();
     }
 
+    @Test(expected = NullPointerException.class)
+    public void testPutFileButFailedToAppend() {
+        DataLakeFileClient fileClient = mock(DataLakeFileClient.class);
+        InputStream stream = mock(InputStream.class);
+        doThrow(NullPointerException.class).when(fileClient).append(any(InputStream.class), anyLong(), anyLong());
+
+        PutAzureDataLakeStorage.uploadContent(fileClient, stream, FILE_DATA.length);
+
+        verify(fileClient).delete();
+    }
+
     private Map<String, String> createAttributesMap() {
         Map<String, String> attributes = new HashMap<>();