embulk-input-gcs 0.2.5 → 0.2.6

data/src/main/java/org/embulk/input/gcs/SingleFileProvider.java

@@ -0,0 +1,143 @@
+package org.embulk.input.gcs;
+
+import com.google.api.client.util.IOUtils;
+import com.google.api.services.storage.Storage;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Throwables;
+import org.embulk.spi.Exec;
+import org.embulk.spi.util.InputStreamFileInput;
+import org.embulk.spi.util.ResumableInputStream;
+import org.embulk.spi.util.RetryExecutor;
+import org.slf4j.Logger;
+
+import java.io.BufferedInputStream;
+import java.io.BufferedOutputStream;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InterruptedIOException;
+import java.util.Iterator;
+
+import static org.embulk.spi.util.RetryExecutor.retryExecutor;
+
+public class SingleFileProvider
+        implements InputStreamFileInput.Provider
+{
+    private final Storage client;
+    private final String bucket;
+    private final Iterator<String> iterator;
+    private final int maxConnectionRetry;
+    private boolean opened = false;
+
+    public SingleFileProvider(PluginTask task, int taskIndex)
+    {
+        this.client = GcsFileInput.newGcsClient(task, GcsFileInput.newGcsAuth(task));
+        this.bucket = task.getBucket();
+        this.iterator = task.getFiles().get(taskIndex).iterator();
+        this.maxConnectionRetry = task.getMaxConnectionRetry();
+    }
+
+    @Override
+    public InputStream openNext() throws IOException
+    {
+        if (opened) {
+            return null;
+        }
+        opened = true;
+        if (!iterator.hasNext()) {
+            return null;
+        }
+        String key = iterator.next();
+        Storage.Objects.Get getObject = client.objects().get(bucket, key);
+        File tempFile = Exec.getTempFileSpace().createTempFile();
+        try (BufferedOutputStream outputStream = new BufferedOutputStream(new FileOutputStream(tempFile))) {
+            IOUtils.copy(getObject.executeMediaAsInputStream(), outputStream);
+        }
+        return new ResumableInputStream(new BufferedInputStream(new FileInputStream(tempFile)), new GcsInputStreamReopener(tempFile, client, bucket, key, maxConnectionRetry));
+    }
+
+    @Override
+    public void close()
+    {
+    }
+
+    @VisibleForTesting
+    static class GcsInputStreamReopener
+            implements ResumableInputStream.Reopener
+    {
+        private final Logger log = Exec.getLogger(GcsInputStreamReopener.class);
+        private final File tempFile;
+        private final Storage client;
+        private final String bucket;
+        private final String key;
+        private final int maxConnectionRetry;
+
+        public GcsInputStreamReopener(File tempFile, Storage client, String bucket, String key, int maxConnectionRetry)
+        {
+            this.tempFile = tempFile;
+            this.client = client;
+            this.bucket = bucket;
+            this.key = key;
+            this.maxConnectionRetry = maxConnectionRetry;
+        }
+
+        @Override
+        public InputStream reopen(final long offset, final Exception closedCause) throws IOException
+        {
+            try {
+                return retryExecutor()
+                        .withRetryLimit(maxConnectionRetry)
+                        .withInitialRetryWait(500)
+                        .withMaxRetryWait(30 * 1000)
+                        .runInterruptible(new RetryExecutor.Retryable<InputStream>() {
+                            @Override
+                            public InputStream call() throws IOException
+                            {
+                                log.warn(String.format("GCS read failed. Retrying GET request with %,d bytes offset", offset), closedCause);
+                                Storage.Objects.Get getObject = client.objects().get(bucket, key);
+
+                                try (BufferedOutputStream outputStream = new BufferedOutputStream(new FileOutputStream(tempFile))) {
+                                    IOUtils.copy(getObject.executeMediaAsInputStream(), outputStream);
+                                }
+                                return new BufferedInputStream(new FileInputStream(tempFile));
+                            }
+
+                            @Override
+                            public boolean isRetryableException(Exception exception)
+                            {
+                                return true;  // TODO
+                            }
+
+                            @Override
+                            public void onRetry(Exception exception, int retryCount, int retryLimit, int retryWait)
+                                    throws RetryExecutor.RetryGiveupException
+                            {
+                                String message = String.format("GCS GET request failed. Retrying %d/%d after %d seconds. Message: %s",
+                                        retryCount, retryLimit, retryWait / 1000, exception.getMessage());
+                                if (retryCount % 3 == 0) {
+                                    log.warn(message, exception);
+                                }
+                                else {
+                                    log.warn(message);
+                                }
+                            }
+
+                            @Override
+                            public void onGiveup(Exception firstException, Exception lastException)
+                                    throws RetryExecutor.RetryGiveupException
+                            {
+                            }
+                        });
+            }
+            catch (RetryExecutor.RetryGiveupException ex) {
+                Throwables.propagateIfInstanceOf(ex.getCause(), IOException.class);
+                throw Throwables.propagate(ex.getCause());
+            }
+            catch (InterruptedException ex) {
+                throw new InterruptedIOException();
+            }
+        }
+    }
+}

data/src/test/java/org/embulk/input/gcs/TestGcsFileInputPlugin.java

@@ -11,7 +11,6 @@ import org.embulk.config.ConfigException;
 import org.embulk.config.ConfigSource;
 import org.embulk.config.TaskReport;
 import org.embulk.config.TaskSource;
-import org.embulk.input.gcs.GcsFileInputPlugin.PluginTask;
 import org.embulk.spi.Exec;
 import org.embulk.spi.FileInputPlugin;
 import org.embulk.spi.FileInputRunner;
@@ -37,7 +36,6 @@ import java.util.Arrays;
 import java.util.List;
 
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
 import static org.junit.Assume.assumeNotNull;
 
 import java.lang.reflect.InvocationTargetException;
@@ -105,7 +103,7 @@ public class TestGcsFileInputPlugin
                 .set("bucket", GCP_BUCKET)
                 .set("path_prefix", "my-prefix");
 
-        GcsFileInputPlugin.PluginTask task = config.loadConfig(PluginTask.class);
+        PluginTask task = config.loadConfig(PluginTask.class);
         assertEquals(true, task.getIncremental());
         assertEquals("private_key", task.getAuthMethod().toString());
         assertEquals("Embulk GCS input plugin", task.getApplicationName());
@@ -124,8 +122,8 @@ public class TestGcsFileInputPlugin
                 .set("p12_keyfile_fullpath", GCP_P12_KEYFILE)
                 .set("parser", parserConfig(schemaConfig()));
 
-        GcsFileInputPlugin.PluginTask task = config.loadConfig(PluginTask.class);
-        assertFalse(task.getFiles().isEmpty());
+        PluginTask task = config.loadConfig(PluginTask.class);
+        assertEquals(2, task.getPathFiles().size());
     }
 
     // both path_prefix and paths are not set
@@ -230,7 +228,7 @@ public class TestGcsFileInputPlugin
 
         Method method = GcsFileInputPlugin.class.getDeclaredMethod("newGcsAuth", PluginTask.class);
         method.setAccessible(true);
-        plugin.newGcsClient(task, (GcsAuthentication) method.invoke(plugin, task)); // no errors happens
+        GcsFileInput.newGcsClient(task, (GcsAuthentication) method.invoke(plugin, task)); // no errors happens
     }
 
     @Test(expected = ConfigException.class)
@@ -251,14 +249,16 @@ public class TestGcsFileInputPlugin
 
         Method method = GcsFileInputPlugin.class.getDeclaredMethod("newGcsAuth", PluginTask.class);
         method.setAccessible(true);
-        plugin.newGcsClient(task, (GcsAuthentication) method.invoke(plugin, task));
+        GcsFileInput.newGcsClient(task, (GcsAuthentication) method.invoke(plugin, task));
     }
 
     @Test
     public void testResume()
     {
         PluginTask task = config.loadConfig(PluginTask.class);
-        task.setFiles(Arrays.asList(new String[]{"in/aa/a"}));
+        FileList.Builder builder = new FileList.Builder(config);
+        builder.add("in/aa/a", 1);
+        task.setFiles(builder.build());
         ConfigDiff configDiff = plugin.resume(task.dump(), 0, new FileInputPlugin.Control()
         {
             @Override
@@ -298,12 +298,44 @@ public class TestGcsFileInputPlugin
 
         Method method = GcsFileInputPlugin.class.getDeclaredMethod("newGcsAuth", PluginTask.class);
         method.setAccessible(true);
-        Storage client = plugin.newGcsClient(task, (GcsAuthentication) method.invoke(plugin, task));
-        List<String> actual = plugin.listGcsFilesByPrefix(client, GCP_BUCKET, GCP_PATH_PREFIX, Optional.<String>absent());
-        assertEquals(expected, actual);
+        Storage client = GcsFileInput.newGcsClient(task, (GcsAuthentication) method.invoke(plugin, task));
+        FileList.Builder builder = new FileList.Builder(config);
+        GcsFileInput.listGcsFilesByPrefix(builder, client, GCP_BUCKET, GCP_PATH_PREFIX, Optional.<String>absent());
+        FileList fileList = builder.build();
+        assertEquals(expected.get(0), fileList.get(0).get(0));
+        assertEquals(expected.get(1), fileList.get(1).get(0));
         assertEquals(GCP_BUCKET_DIRECTORY + "sample_02.csv", configDiff.get(String.class, "last_path"));
     }
 
+    @Test
+    public void testListFilesByPrefixWithPattern()
+            throws NoSuchMethodException, IllegalAccessException, InvocationTargetException
+    {
+        List<String> expected = Arrays.asList(
+                GCP_BUCKET_DIRECTORY + "sample_01.csv"
+        );
+
+        ConfigSource configWithPattern = config.deepCopy().set("path_match_pattern", "1");
+        PluginTask task = configWithPattern.loadConfig(PluginTask.class);
+        ConfigDiff configDiff = plugin.transaction(configWithPattern, new FileInputPlugin.Control() {
+            @Override
+            public List<TaskReport> run(TaskSource taskSource, int taskCount)
+            {
+                assertEquals(1, taskCount);
+                return emptyTaskReports(taskCount);
+            }
+        });
+
+        Method method = GcsFileInputPlugin.class.getDeclaredMethod("newGcsAuth", PluginTask.class);
+        method.setAccessible(true);
+        Storage client = GcsFileInput.newGcsClient(task, (GcsAuthentication) method.invoke(plugin, task));
+        FileList.Builder builder = new FileList.Builder(configWithPattern);
+        GcsFileInput.listGcsFilesByPrefix(builder, client, GCP_BUCKET, GCP_PATH_PREFIX, Optional.<String>absent());
+        FileList fileList = builder.build();
+        assertEquals(expected.get(0), fileList.get(0).get(0));
+        assertEquals(GCP_BUCKET_DIRECTORY + "sample_01.csv", configDiff.get(String.class, "last_path"));
+    }
+
     @Test
     public void testListFilesByPrefixIncrementalFalse() throws Exception
     {
@@ -324,8 +356,9 @@ public class TestGcsFileInputPlugin
 
         Method method = GcsFileInputPlugin.class.getDeclaredMethod("newGcsAuth", PluginTask.class);
         method.setAccessible(true);
-        Storage client = plugin.newGcsClient(task, (GcsAuthentication) method.invoke(plugin, task));
-        plugin.listGcsFilesByPrefix(client, "non-exists-bucket", "prefix", Optional.<String>absent()); // no errors happens
+        Storage client = GcsFileInput.newGcsClient(task, (GcsAuthentication) method.invoke(plugin, task));
+        FileList.Builder builder = new FileList.Builder(config);
+        GcsFileInput.listGcsFilesByPrefix(builder, client, "non-exists-bucket", "prefix", Optional.<String>absent()); // no errors happens
     }
 
     @Test
@@ -379,10 +412,10 @@ public class TestGcsFileInputPlugin
         PluginTask task = config.loadConfig(PluginTask.class);
         runner.transaction(config, new Control());
 
-        Method method = GcsFileInputPlugin.class.getDeclaredMethod("newGcsAuth", PluginTask.class);
+        Method method = GcsFileInput.class.getDeclaredMethod("newGcsAuth", PluginTask.class);
         method.setAccessible(true);
-        Storage client = plugin.newGcsClient(task, (GcsAuthentication) method.invoke(plugin, task));
-        task.setFiles(plugin.listFiles(task, client));
+        Storage client = GcsFileInput.newGcsClient(task, (GcsAuthentication) method.invoke(plugin, task));
+        task.setFiles(GcsFileInput.listFiles(task, client));
 
         assertRecords(config, output);
     }
@@ -402,14 +435,14 @@ public class TestGcsFileInputPlugin
         PluginTask task = config.loadConfig(PluginTask.class);
         runner.transaction(config, new Control());
 
-        Method method = GcsFileInputPlugin.class.getDeclaredMethod("newGcsAuth", PluginTask.class);
+        Method method = GcsFileInput.class.getDeclaredMethod("newGcsAuth", PluginTask.class);
         method.setAccessible(true);
-        Storage client = plugin.newGcsClient(task, (GcsAuthentication) method.invoke(plugin, task));
+        Storage client = GcsFileInput.newGcsClient(task, (GcsAuthentication) method.invoke(plugin, task));
         File tempFile = Exec.getTempFileSpace().createTempFile();
-        task.setFiles(plugin.listFiles(task, client));
+        task.setFiles(GcsFileInput.listFiles(task, client));
 
         String key = GCP_BUCKET_DIRECTORY + "sample_01.csv";
-        GcsFileInputPlugin.GcsInputStreamReopener opener = new GcsFileInputPlugin.GcsInputStreamReopener(tempFile, client, GCP_BUCKET, key, MAX_CONNECTION_RETRY);
+        SingleFileProvider.GcsInputStreamReopener opener = new SingleFileProvider.GcsInputStreamReopener(tempFile, client, GCP_BUCKET, key, MAX_CONNECTION_RETRY);
         try (InputStream in = opener.reopen(0, new RuntimeException())) {
             BufferedReader r = new BufferedReader(new InputStreamReader(in));
             assertEquals("id,account,time,purchase,comment", r.readLine());
@@ -420,7 +453,7 @@ public class TestGcsFileInputPlugin
     public void testBase64()
             throws NoSuchMethodException, IllegalAccessException, InvocationTargetException
     {
-        Method method = GcsFileInputPlugin.class.getDeclaredMethod("base64Encode", String.class);
+        Method method = GcsFileInput.class.getDeclaredMethod("base64Encode", String.class);
         method.setAccessible(true);
 
         assertEquals("CgFj", method.invoke(plugin, "c"));

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: embulk-input-gcs
 version: !ruby/object:Gem::Version
-  version: 0.2.5
+  version: 0.2.6
 platform: ruby
 authors:
 - Satoshi Akama
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2017-05-19 00:00:00.000000000 Z
+date: 2018-03-05 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   requirement: !ruby/object:Gem::Requirement
@@ -58,8 +58,12 @@ files:
 - gradlew.bat
 - lib/embulk/input/gcs.rb
 - settings.gradle
+- src/main/java/org/embulk/input/gcs/FileList.java
 - src/main/java/org/embulk/input/gcs/GcsAuthentication.java
+- src/main/java/org/embulk/input/gcs/GcsFileInput.java
 - src/main/java/org/embulk/input/gcs/GcsFileInputPlugin.java
+- src/main/java/org/embulk/input/gcs/PluginTask.java
+- src/main/java/org/embulk/input/gcs/SingleFileProvider.java
 - src/test/java/org/embulk/input/gcs/TestGcsAuthentication.java
 - src/test/java/org/embulk/input/gcs/TestGcsFileInputPlugin.java
 - src/test/resources/sample_01.csv
@@ -67,7 +71,7 @@ files:
 - src/test/resources/secretkeys.tar.enc
 - classpath/commons-codec-1.3.jar
 - classpath/commons-logging-1.1.1.jar
-- classpath/embulk-input-gcs-0.2.5.jar
+- classpath/embulk-input-gcs-0.2.6.jar
 - classpath/google-api-client-1.21.0.jar
 - classpath/google-api-services-storage-v1-rev59-1.21.0.jar
 - classpath/google-http-client-1.21.0.jar