Can choose a policy in delete_in_advance option

Author: civitaspo

README.md

@@ -23,10 +23,13 @@ A File Output Plugin for Embulk to write HDFS.
 - **overwrite** overwrite files when the same filenames already exists (boolean, default: `false`)
     - *caution*: even if this property is `true`, this does not mean ensuring the idempotence. if you want to ensure the idempotence, you need the procedures to remove output files after or before running. 
 - **doas** username which access to Hdfs (string, default: executed user)
-- **delete_in_advance** delete files and directories having `path_prefix` in advance (boolean, default: `false`)
+- **delete_in_advance** delete files and directories having `path_prefix` in advance (enum, default: `NONE`)
+    - `NONE`: do nothing
+    - `FILE_ONLY`: delete files
+    - `RECURSIVE`: delete files and directories
 
 ## CAUTION
-If you use `hadoop` user (hdfs admin user) as `doas`, and if `delete_in_advance` is true,
+If you use `hadoop` user (hdfs admin user) as `doas`, and if `delete_in_advance` is `RECURSIVE`,
 `embulk-output-hdfs` can delete any files and directories you indicate as `path_prefix`,
 this means `embulk-output-hdfs` can destroy your hdfs.
 So, please be careful when you use `delete_in_advance` option and `doas` option ...

src/main/java/org/embulk/output/hdfs/HdfsFileOutputPlugin.java

@@ -9,6 +9,7 @@
 import org.embulk.config.Config;
 import org.embulk.config.ConfigDefault;
 import org.embulk.config.ConfigDiff;
+import org.embulk.config.ConfigException;
 import org.embulk.config.ConfigSource;
 import org.embulk.config.Task;
 import org.embulk.config.TaskReport;
@@ -28,6 +29,8 @@
 import java.util.List;
 import java.util.Map;
 
+import static org.embulk.output.hdfs.HdfsFileOutputPlugin.PluginTask.*;
+
 public class HdfsFileOutputPlugin
         implements FileOutputPlugin
 {
@@ -48,7 +51,7 @@ public interface PluginTask
         String getPathPrefix();
 
         @Config("file_ext")
-        String getFileNameExtension();
+        String getFileExt();
 
         @Config("sequence_format")
         @ConfigDefault("\"%03d.%02d.\"")
@@ -66,9 +69,10 @@ public interface PluginTask
         @ConfigDefault("null")
         Optional<String> getDoas();
 
+        enum DeleteInAdvancePolicy{ NONE, FILE_ONLY, RECURSIVE}
         @Config("delete_in_advance")
-        @ConfigDefault("false")
-        boolean getDeleteInAdvance();
+        @ConfigDefault("\"NONE\"")
+        DeleteInAdvancePolicy getDeleteInAdvance();
     }
 
     @Override
@@ -77,19 +81,13 @@ public ConfigDiff transaction(ConfigSource config, int taskCount,
     {
         PluginTask task = config.loadConfig(PluginTask.class);
 
-        if (task.getDeleteInAdvance()) {
-            final String pathPrefix = strftime(task.getPathPrefix(), task.getRewindSeconds());
-            final Path globPath = new Path(pathPrefix + "*");
-            try {
-                FileSystem fs = getFs(task);
-                for (FileStatus status : fs.globStatus(globPath)) {
-                    logger.debug("delete in advance: {}", status.getPath());
-                    fs.delete(status.getPath(), true);
-                }
-            }
-            catch (IOException e) {
-                throw Throwables.propagate(e);
-            }
+        try {
+            String pathPrefix = strftime(task.getPathPrefix(), task.getRewindSeconds());
+            FileSystem fs = getFs(task);
+            deleteInAdvance(fs, pathPrefix, task.getDeleteInAdvance());
+        }
+        catch (IOException e) {
+            throw Throwables.propagate(e);
         }
 
         control.run(task.dump());
@@ -117,7 +115,7 @@ public TransactionalFileOutput open(TaskSource taskSource, final int taskIndex)
         final PluginTask task = taskSource.loadTask(PluginTask.class);
 
         final String pathPrefix = strftime(task.getPathPrefix(), task.getRewindSeconds());
-        final String pathSuffix = task.getFileNameExtension();
+        final String pathSuffix = task.getFileExt();
         final String sequenceFormat = task.getSequenceFormat();
 
         return new TransactionalFileOutput()
@@ -231,4 +229,31 @@ private String strftime(final String raw, final int rewind_seconds)
                 String.format("(Time.now - %s).strftime('%s')", String.valueOf(rewind_seconds), raw));
         return resolved.toString();
     }
+
+    private void deleteInAdvance(FileSystem fs, String pathPrefix, DeleteInAdvancePolicy deleteInAdvancePolicy)
+            throws IOException
+    {
+        final Path globPath = new Path(pathPrefix + "*");
+        switch (deleteInAdvancePolicy) {
+            case NONE:
+                // do nothing
+                break;
+            case FILE_ONLY:
+                for (FileStatus status : fs.globStatus(globPath)) {
+                    if (status.isFile()) {
+                        logger.debug("delete in advance: {}", status.getPath());
+                        fs.delete(status.getPath(), false);
+                    }
+                }
+                break;
+            case RECURSIVE:
+                for (FileStatus status : fs.globStatus(globPath)) {
+                    logger.debug("delete in advance: {}", status.getPath());
+                    fs.delete(status.getPath(), true);
+                }
+                break;
+            default:
+                throw new ConfigException("`delete_in_advance` must not null.");
+        }
+    }
 }

src/test/java/org/embulk/output/hdfs/TestHdfsFileOutputPlugin.java

@@ -102,14 +102,14 @@ public void testDefaultValues()
         ConfigSource config = getBaseConfigSource();
         PluginTask task = config.loadConfig(PluginTask.class);
         assertEquals(pathPrefix, task.getPathPrefix());
-        assertEquals("csv", task.getFileNameExtension());
+        assertEquals("csv", task.getFileExt());
         assertEquals("%03d.%02d.", task.getSequenceFormat());
         assertEquals(Lists.newArrayList(), task.getConfigFiles());
         assertEquals(Maps.newHashMap(), task.getConfig());
         assertEquals(0, task.getRewindSeconds());
         assertEquals(false, task.getOverwrite());
         assertEquals(Optional.absent(), task.getDoas());
-        assertEquals(false, task.getDeleteInAdvance());
+        assertEquals(PluginTask.DeleteInAdvancePolicy.NONE, task.getDeleteInAdvance());
     }
 
     @Test(expected = ConfigException.class)
@@ -119,24 +119,25 @@ public void testRequiredValues()
         PluginTask task = config.loadConfig(PluginTask.class);
     }
 
-    private List<String> lsR(List<String> fileNames, java.nio.file.Path dir)
+    private List<String> lsR(List<String> names, java.nio.file.Path dir)
     {
         try (DirectoryStream<java.nio.file.Path> stream = Files.newDirectoryStream(dir)) {
             for (java.nio.file.Path path : stream) {
                 if (path.toFile().isDirectory()) {
                     logger.debug("[lsR] find a directory: {}", path.toAbsolutePath().toString());
-                    lsR(fileNames, path);
+                    names.add(path.toAbsolutePath().toString());
+                    lsR(names, path);
                 }
                 else {
                     logger.debug("[lsR] find a file: {}", path.toAbsolutePath().toString());
-                    fileNames.add(path.toAbsolutePath().toString());
+                    names.add(path.toAbsolutePath().toString());
                 }
             }
         }
         catch (IOException e) {
             logger.debug(e.getMessage(), e);
         }
-        return fileNames;
+        return names;
     }
 
     private void run(ConfigSource config)
@@ -218,12 +219,12 @@ public void testBulkLoad()
     }
 
     @Test
-    public void testDeleteInAdvance()
+    public void testDeleteRECURSIVEInAdvance()
             throws IOException
     {
         for (int n = 0; n <= 10; n++) {
-            tmpFolder.newFile("embulk-output-hdfs_" + n + ".txt");
-            tmpFolder.newFolder("embulk-output-hdfs_" + n);
+            tmpFolder.newFile("embulk-output-hdfs_file_" + n + ".txt");
+            tmpFolder.newFolder("embulk-output-hdfs_directory_" + n);
         }
 
         List<String> fileListBeforeRun = lsR(Lists.<String>newArrayList(), Paths.get(tmpFolder.getRoot().getAbsolutePath()));
@@ -233,16 +234,47 @@ public void testDeleteInAdvance()
                         .set("fs.hdfs.impl", "org.apache.hadoop.fs.RawLocalFileSystem")
                         .set("fs.file.impl", "org.apache.hadoop.fs.RawLocalFileSystem")
                         .set("fs.defaultFS", "file:///"))
-                .set("delete_in_advance", true);
+                .set("delete_in_advance", "RECURSIVE");
 
         run(config);
 
         List<String> fileListAfterRun = lsR(Lists.<String>newArrayList(), Paths.get(tmpFolder.getRoot().getAbsolutePath()));
         assertNotEquals(fileListBeforeRun, fileListAfterRun);
+        assertThat(fileListAfterRun, not(hasItem(containsString("embulk-output-hdfs_directory"))));
         assertThat(fileListAfterRun, not(hasItem(containsString("txt"))));
         assertThat(fileListAfterRun, hasItem(containsString(pathPrefix + "001.00.csv")));
         assertRecordsInFile(String.format("%s/%s001.00.csv",
                 tmpFolder.getRoot().getAbsolutePath(),
                 pathPrefix));
     }
+
+    @Test
+    public void testDeleteFILE_ONLYInAdvance()
+            throws IOException
+    {
+        for (int n = 0; n <= 10; n++) {
+            tmpFolder.newFile("embulk-output-hdfs_file_" + n + ".txt");
+            tmpFolder.newFolder("embulk-output-hdfs_directory_" + n);
+        }
+
+        List<String> fileListBeforeRun = lsR(Lists.<String>newArrayList(), Paths.get(tmpFolder.getRoot().getAbsolutePath()));
+
+        ConfigSource config = getBaseConfigSource()
+                .setNested("config", Exec.newConfigSource()
+                        .set("fs.hdfs.impl", "org.apache.hadoop.fs.RawLocalFileSystem")
+                        .set("fs.file.impl", "org.apache.hadoop.fs.RawLocalFileSystem")
+                        .set("fs.defaultFS", "file:///"))
+                .set("delete_in_advance", "FILE_ONLY");
+
+        run(config);
+
+        List<String> fileListAfterRun = lsR(Lists.<String>newArrayList(), Paths.get(tmpFolder.getRoot().getAbsolutePath()));
+        assertNotEquals(fileListBeforeRun, fileListAfterRun);
+        assertThat(fileListAfterRun, not(hasItem(containsString("txt"))));
+        assertThat(fileListAfterRun, hasItem(containsString("embulk-output-hdfs_directory")));
+        assertThat(fileListAfterRun, hasItem(containsString(pathPrefix + "001.00.csv")));
+        assertRecordsInFile(String.format("%s/%s001.00.csv",
+                tmpFolder.getRoot().getAbsolutePath(),
+                pathPrefix));
+    }
 }