delete files in advance

kekekenta · kekekenta · commit 1b7e0a9b0e66 · 2020-03-06T15:52:55.000+09:00
diff --git a/README.md b/README.md
@@ -24,6 +24,7 @@ Google Cloud Storage output plugin for [Embulk](https://github.com/embulk/embulk
 - **json_keyfile** fullpath of json_key (string, required when auth_method is json_key)
 - **application_name**: Application name, anything you like (string, optional, default value is "embulk-output-gcs")
 - **max_connection_retry**: Number of connection retries to GCS (number, default value is 10)
+- **delete_in_advance**: Delete Bucket/Prefix matched files in advance (boolean, default value is false)
 
 ## Example
 
diff --git a/src/main/java/org/embulk/output/GcsOutputPlugin.java b/src/main/java/org/embulk/output/GcsOutputPlugin.java
@@ -1,6 +1,8 @@
 package org.embulk.output;
 
 import com.google.api.services.storage.Storage;
+import com.google.api.services.storage.model.StorageObject;
+import com.google.api.services.storage.model.Objects;
 import com.google.common.base.Throwables;
 import org.embulk.config.ConfigDiff;
 import org.embulk.config.ConfigException;
@@ -11,15 +13,25 @@
 import org.embulk.spi.FileOutputPlugin;
 import org.embulk.spi.TransactionalFileOutput;
 import org.embulk.spi.unit.LocalFile;
+import org.embulk.spi.util.RetryExecutor.RetryGiveupException;
+import org.embulk.spi.util.RetryExecutor.Retryable;
+import org.slf4j.Logger;
 
 import java.io.IOException;
+import java.io.InterruptedIOException;
 import java.security.GeneralSecurityException;
+import java.util.ArrayList;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Optional;
 import java.util.function.Function;
 
+import static org.embulk.spi.util.RetryExecutor.retryExecutor;
+
 public class GcsOutputPlugin implements FileOutputPlugin
 {
+    private static final Logger logger = Exec.getLogger(GcsOutputPlugin.class);
+
     @Override
     public ConfigDiff transaction(ConfigSource config,
                                   int taskCount,
@@ -50,6 +62,10 @@ else if (task.getAuthMethod().getString().equals("private_key")) {
             }
         }
 
+        if (task.getDeleteInAdvance()) {
+            deleteFiles(task);
+        }
+
         return resume(task.dump(), taskCount, control);
     }
 
@@ -78,6 +94,133 @@ public TransactionalFileOutput open(TaskSource taskSource, final int taskIndex)
         return new GcsTransactionalFileOutput(task, client, taskIndex);
     }
 
+    private void deleteFiles(PluginTask task)
+    {
+        logger.info("Start delete files operation");
+        Storage client = createClient(task);
+        try {
+            List<StorageObject> items = listObjectsWithRetry(client, task.getBucket(), task.getPathPrefix(), task.getMaxConnectionRetry());
+            if (items.size() == 0) {
+                logger.info("no files were found");
+                return;
+            }
+            for (StorageObject item : items) {
+                deleteObjectWithRetry(client, item, task.getMaxConnectionRetry());
+                logger.info("delete file: {}/{}", item.getBucket(), item.getName());
+            }
+        }
+        catch (IOException ex) {
+            throw new ConfigException(ex);
+        }
+    }
+
+    private List<StorageObject> listObjectsWithRetry(Storage client, String bucket, String prefix, int maxConnectionRetry) throws IOException
+    {
+        try {
+            return retryExecutor()
+                .withRetryLimit(maxConnectionRetry)
+                .withInitialRetryWait(500)
+                .withMaxRetryWait(30 * 1000)
+                .runInterruptible(new Retryable<List<StorageObject>>() {
+                @Override
+                public List<StorageObject> call() throws IOException
+                {
+                    Storage.Objects.List listObjects = client.objects().list(bucket).setDelimiter("/").setPrefix(prefix);
+                    List<StorageObject> items = new LinkedList<StorageObject>();
+                    String token = null;
+                    do {
+                        Objects objects = listObjects.execute();
+                        if (objects.getItems() == null) {
+                            break;
+                        }
+                        items.addAll(objects.getItems());
+                        token = objects.getNextPageToken();
+                        listObjects.setPageToken(token);
+                    } while (token != null);
+                    return items;
+                }
+
+                @Override
+                public boolean isRetryableException(Exception exception)
+                {
+                    return true;
+                }
+
+                @Override
+                public void onRetry(Exception exception, int retryCount, int retryLimit, int retryWait) throws RetryGiveupException
+                {
+                    String message = String.format("GCS list request failed. Retrying %d/%d after %d seconds. Message: %s: %s",
+                                    retryCount, retryLimit, retryWait / 1000, exception.getClass(), exception.getMessage());
+                    if (retryCount % 3 == 0) {
+                        logger.warn(message, exception);
+                    }
+                    else {
+                        logger.warn(message);
+                    }
+                }
+
+                @Override
+                public void onGiveup(Exception firstException, Exception lastException) throws RetryGiveupException
+                {
+                }
+            });
+        }
+        catch (RetryGiveupException ex) {
+            throw Throwables.propagate(ex.getCause());
+        }
+        catch (InterruptedException ex) {
+            throw new InterruptedIOException();
+        }
+    }
+
+    private Void deleteObjectWithRetry(Storage client, StorageObject item, int maxConnectionRetry) throws IOException
+    {
+        try {
+            return retryExecutor()
+                .withRetryLimit(maxConnectionRetry)
+                .withInitialRetryWait(500)
+                .withMaxRetryWait(30 * 1000)
+                .runInterruptible(new Retryable<Void>() {
+                @Override
+                public Void call() throws IOException
+                {
+                    client.objects().delete(item.getBucket(), item.getName()).execute();
+                    return null;
+                }
+
+                @Override
+                public boolean isRetryableException(Exception exception)
+                {
+                    return true;
+                }
+
+                @Override
+                public void onRetry(Exception exception, int retryCount, int retryLimit, int retryWait) throws RetryGiveupException
+                {
+                    String message = String.format("GCS delete request failed. Retrying %d/%d after %d seconds. Message: %s: %s",
+                                    retryCount, retryLimit, retryWait / 1000, exception.getClass(), exception.getMessage());
+                    if (retryCount % 3 == 0) {
+                        logger.warn(message, exception);
+                    }
+                    else {
+                        logger.warn(message);
+                    }
+                }
+
+                @Override
+                public void onGiveup(Exception firstException, Exception lastException) throws RetryGiveupException
+                {
+                }
+            });
+        }
+        catch (RetryGiveupException ex) {
+            throw Throwables.propagate(ex.getCause());
+        }
+        catch (InterruptedException ex) {
+            throw new InterruptedIOException();
+        }
+    }
+
     private GcsAuthentication newGcsAuth(PluginTask task)
     {
         try {
diff --git a/src/main/java/org/embulk/output/PluginTask.java b/src/main/java/org/embulk/output/PluginTask.java
@@ -55,4 +55,8 @@ public interface PluginTask extends Task
     @Config("max_connection_retry")
     @ConfigDefault("10") // 10 times retry to connect GCS server if failed.
     int getMaxConnectionRetry();
+
+    @Config("delete_in_advance")
+    @ConfigDefault("false")
+    boolean getDeleteInAdvance();
 }
diff --git a/src/test/java/org/embulk/output/TestGcsOutputPlugin.java b/src/test/java/org/embulk/output/TestGcsOutputPlugin.java
@@ -1,6 +1,9 @@
 package org.embulk.output;
 
+import com.google.api.client.http.InputStreamContent;
 import com.google.api.services.storage.Storage;
+import com.google.api.services.storage.Storage.Objects.Get;
+import com.google.api.services.storage.model.StorageObject;
 import com.google.common.base.Optional;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -28,7 +31,9 @@
 import static org.junit.Assume.assumeNotNull;
 
 import java.io.BufferedReader;
+import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
+import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStream;
@@ -99,6 +104,7 @@ public void checkDefaultValues()
 
         PluginTask task = config.loadConfig(PluginTask.class);
         assertEquals("private_key", task.getAuthMethod().toString());
+        assertEquals(false, task.getDeleteInAdvance());
     }
 
     // p12_keyfile is null when auth_method is private_key
@@ -296,6 +302,30 @@ public void testGenerateRemotePath() throws Exception
         assertEquals("sample.000.01.csv", method.invoke(plugin, "......///sample", task.getSequenceFormat(), 0, 1, ".csv"));
     }
 
+    @Test
+    public void testDeleteFilesInAdvanceSuccessfully() throws Exception
+    {
+        ConfigSource configSource = config();
+        PluginTask task = configSource.loadConfig(PluginTask.class);
+
+        Method createClientMethod = GcsOutputPlugin.class.getDeclaredMethod("createClient", PluginTask.class);
+        createClientMethod.setAccessible(true);
+        Storage client = (Storage) createClientMethod.invoke(plugin, task);
+
+        uploadEmptyFile(client, task.getBucket(), task.getPathPrefix() + ".001.csv");
+        uploadEmptyFile(client, task.getBucket(), task.getPathPrefix() + ".002.csv");
+
+        assertEquals(true, isFileExist(client, task.getBucket(), task.getPathPrefix() + ".001.csv"));
+        assertEquals(true, isFileExist(client, task.getBucket(), task.getPathPrefix() + ".002.csv"));
+
+        Method deleteFilesMethod = GcsOutputPlugin.class.getDeclaredMethod("deleteFiles", PluginTask.class);
+        deleteFilesMethod.setAccessible(true);
+        deleteFilesMethod.invoke(plugin, task);
+
+        assertEquals(false, isFileExist(client, task.getBucket(), task.getPathPrefix() + ".001.csv"));
+        assertEquals(false, isFileExist(client, task.getBucket(), task.getPathPrefix() + ".002.csv"));
+    }
+
     public ConfigSource config()
     {
         return Exec.newConfigSource()
@@ -416,6 +446,23 @@ private ImmutableList<List<String>> getFileContentsFromGcs(String path) throws E
         return builder.build();
     }
 
+    private static void uploadEmptyFile(Storage client, String gcsBucket, String gcsPath) throws IOException
+    {
+        InputStreamContent content = new InputStreamContent("text/plain", new ByteArrayInputStream(new byte[0]));
+        StorageObject object = new StorageObject().setName(gcsPath);
+        client.objects().insert(gcsBucket, object, content).execute();
+    }
+
+    private static boolean isFileExist(Storage client, String gcsBucket, String gcsPath)
+    {
+        try {
+            client.objects().get(gcsBucket, gcsPath).setAlt("json").execute();
+        } catch (IOException ex) {
+            return false;
+        }
+        return true;
+    }
+
     private static String getDirectory(String dir)
     {
         if (dir != null && !dir.endsWith("/")) {

Original file line number	Diff line number	Diff line change
`@@ -55,4 +55,8 @@ public interface PluginTask extends Task`
`55`	`55`	`@Config("max_connection_retry")`
`56`	`56`	`@ConfigDefault("10") // 10 times retry to connect GCS server if failed.`
`57`	`57`	`int getMaxConnectionRetry();`
	`58`	`+`
	`59`	`+ @Config("delete_in_advance")`
	`60`	`+ @ConfigDefault("false")`
	`61`	`+ boolean getDeleteInAdvance();`
`58`	`62`	`}`