embulk-input-gcs 0.3.0 → 0.3.1

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

@@ -1,34 +1,26 @@
 package org.embulk.input.gcs;
 
-import com.google.api.client.http.HttpResponseException;
-import com.google.api.services.storage.Storage;
-import com.google.api.services.storage.model.Bucket;
-import com.google.api.services.storage.model.Objects;
-import com.google.api.services.storage.model.StorageObject;
+import com.google.api.gax.paging.Page;
+import com.google.cloud.storage.Blob;
+import com.google.cloud.storage.Storage;
+import com.google.cloud.storage.StorageException;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Charsets;
 import com.google.common.io.BaseEncoding;
 import org.embulk.config.ConfigException;
 import org.embulk.config.TaskReport;
 import org.embulk.spi.Exec;
 import org.embulk.spi.TransactionalFileInput;
-import org.embulk.spi.unit.LocalFile;
 import org.embulk.spi.util.InputStreamFileInput;
 import org.slf4j.Logger;
 
-import java.io.IOException;
-import java.math.BigInteger;
-import java.security.GeneralSecurityException;
-import java.util.List;
-import java.util.Optional;
-import java.util.function.Function;
-
 public class GcsFileInput
         extends InputStreamFileInput
         implements TransactionalFileInput
 {
-    private static final Logger log = Exec.getLogger(org.embulk.input.gcs.GcsFileInput.class);
+    private static final Logger LOG = Exec.getLogger(org.embulk.input.gcs.GcsFileInput.class);
 
-    public GcsFileInput(PluginTask task, int taskIndex)
+    GcsFileInput(PluginTask task, int taskIndex)
     {
         super(task.getBufferAllocator(), new SingleFileProvider(task, taskIndex));
     }
@@ -47,149 +39,77 @@ public class GcsFileInput
     {
     }
 
-    public static GcsAuthentication newGcsAuth(PluginTask task)
-    {
-        try {
-            return new GcsAuthentication(
-                    task.getAuthMethod().getString(),
-                    task.getServiceAccountEmail(),
-                    task.getP12Keyfile().map(localFileToPathString()),
-                    task.getJsonKeyfile().map(localFileToPathString()),
-                    task.getApplicationName()
-            );
-        }
-        catch (GeneralSecurityException | IOException ex) {
-            throw new ConfigException(ex);
-        }
-    }
-
-    protected static Storage newGcsClient(final PluginTask task, final GcsAuthentication auth)
-    {
-        Storage client = null;
-        try {
-            client = auth.getGcsClient(task.getBucket(), task.getMaxConnectionRetry());
-        }
-        catch (IOException ex) {
-            throw new ConfigException(ex);
-        }
-
-        return client;
-    }
-
-    private static Function<LocalFile, String> localFileToPathString()
-    {
-        return new Function<LocalFile, String>()
-        {
-            public String apply(LocalFile file)
-            {
-                return file.getPath().toString();
-            }
-        };
-    }
-
-    public static FileList listFiles(PluginTask task, Storage client)
-    {
-        String bucket = task.getBucket();
-
-        FileList.Builder builder = new FileList.Builder(task);
-        listGcsFilesByPrefix(builder, client, bucket, task.getPathPrefix().get(), task.getLastPath());
-        return builder.build();
-    }
-
     /**
      * Lists GCS filenames filtered by prefix.
      *
      * The resulting list does not include the file that's size == 0.
      */
-    public static void listGcsFilesByPrefix(FileList.Builder builder, Storage client, String bucket,
-                                             String prefix, Optional<String> lastPath)
+    static FileList listFiles(PluginTask task)
     {
-        String lastKey = lastPath.isPresent() ? base64Encode(lastPath.get()) : null;
-
-        // @see https://cloud.google.com/storage/docs/json_api/v1/objects#resource
-        if (log.isDebugEnabled()) {
-            try {
-                Storage.Buckets.Get getBucket = client.buckets().get(bucket);
-                getBucket.setProjection("full");
-                Bucket bk = getBucket.execute();
+        Storage client = AuthUtils.newClient(task);
+        String bucket = task.getBucket();
 
-                log.debug("bucket name: " + bucket);
-                log.debug("bucket location: " + bk.getLocation());
-                log.debug("bucket timeCreated: " + bk.getTimeCreated());
-                log.debug("bucket owner: " + bk.getOwner());
-            }
-            catch (IOException e) {
-                log.warn("Could not access to bucket:" + bucket);
-                log.warn(e.getMessage());
-            }
+        // @see https://cloud.google.com/storage/docs/json_api/v1/buckets/get
+        if (LOG.isDebugEnabled()) {
+            printBucketInfo(client, bucket);
         }
 
+        String prefix = task.getPathPrefix().orElse("");
+        String lastKey = task.getLastPath().isPresent() ? base64Encode(task.getLastPath().get()) : "";
+        FileList.Builder builder = new FileList.Builder(task);
+
         try {
             // @see https://cloud.google.com/storage/docs/json_api/v1/objects/list
-            Storage.Objects.List listObjects = client.objects().list(bucket);
-            listObjects.setPrefix(prefix);
-            listObjects.setPageToken(lastKey);
-            do {
-                Objects objects = listObjects.execute();
-                List<StorageObject> items = objects.getItems();
-                if (items == null) {
-                    log.info(String.format("No file was found in bucket:%s prefix:%s", bucket, prefix));
-                    break;
-                }
-                for (StorageObject o : items) {
-                    if (o.getSize().compareTo(BigInteger.ZERO) > 0) {
-                        builder.add(o.getName(), o.getSize().longValue());
-                    }
-                    log.debug("filename: " + o.getName());
-                    log.debug("updated: " + o.getUpdated());
+            Page<Blob> blobs = client.list(bucket, Storage.BlobListOption.prefix(prefix), Storage.BlobListOption.pageToken(lastKey));
+            for (Blob blob : blobs.iterateAll()) {
+                if (blob.getSize() > 0) {
+                    builder.add(blob.getName(), blob.getSize());
                 }
-                lastKey = objects.getNextPageToken();
-                listObjects.setPageToken(lastKey);
-            } while (lastKey != null);
+                LOG.debug("filename: {}", blob.getName());
+                LOG.debug("updated: {}", blob.getUpdateTime());
+            }
         }
-        catch (IOException e) {
-            if ((e instanceof HttpResponseException) && ((HttpResponseException) e).getStatusCode() == 400) {
+        catch (RuntimeException e) {
+            if ((e instanceof StorageException) && ((StorageException) e).getCode() == 400) {
                 throw new ConfigException(String.format("Files listing failed: bucket:%s, prefix:%s, last_path:%s", bucket, prefix, lastKey), e);
             }
 
-            log.warn(String.format("Could not get file list from bucket:%s", bucket));
-            log.warn(e.getMessage());
+            LOG.warn(String.format("Could not get file list from bucket:%s", bucket));
+            LOG.warn(e.getMessage());
         }
+        return builder.build();
     }
 
     // String nextToken = base64Encode(0x0a + 0x01~0x27 + filePath);
-    private static String base64Encode(String path)
+    @VisibleForTesting
+    static String base64Encode(String path)
     {
         byte[] encoding;
         byte[] utf8 = path.getBytes(Charsets.UTF_8);
-        log.debug(String.format("path string: %s ,path length:%s \" + ", path, utf8.length));
+        LOG.debug("path string: {} ,path length:{} \" + ", path, utf8.length);
 
         encoding = new byte[utf8.length + 2];
         encoding[0] = 0x0a;
-        encoding[1] = new Byte(String.valueOf(path.length()));
+        encoding[1] = Byte.valueOf(String.valueOf(path.length()));
         System.arraycopy(utf8, 0, encoding, 2, utf8.length);
 
         String s = BaseEncoding.base64().encode(encoding);
-        log.debug(String.format("last_path(base64 encoded): %s", s));
+        LOG.debug("last_path(base64 encoded): {}", s);
         return s;
     }
 
-    public enum AuthMethod
+    private static void printBucketInfo(Storage client, String bucket)
     {
-        private_key("private_key"),
-        compute_engine("compute_engine"),
-        json_key("json_key");
-
-        private final String string;
-
-        AuthMethod(String string)
-        {
-            this.string = string;
-        }
-
-        public String getString()
-        {
-            return string;
-        }
+        // get Bucket
+        Storage.BucketGetOption fields = Storage.BucketGetOption.fields(
+                Storage.BucketField.LOCATION,
+                Storage.BucketField.TIME_CREATED,
+                Storage.BucketField.OWNER
+        );
+        com.google.cloud.storage.Bucket bk = client.get(bucket, fields);
+        LOG.debug("bucket name: {}", bk.getName());
+        LOG.debug("bucket location: {}", bk.getLocation());
+        LOG.debug("bucket timeCreated: {}", bk.getCreateTime());
+        LOG.debug("bucket owner: {}", bk.getOwner());
     }
 }

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

@@ -1,6 +1,5 @@
 package org.embulk.input.gcs;
 
-import com.google.api.services.storage.Storage;
 import com.google.common.base.Throwables;
 import org.embulk.config.ConfigDiff;
 import org.embulk.config.ConfigException;
@@ -11,19 +10,14 @@ import org.embulk.spi.Exec;
 import org.embulk.spi.FileInputPlugin;
 import org.embulk.spi.TransactionalFileInput;
 import org.embulk.spi.unit.LocalFile;
-import org.slf4j.Logger;
 
 import java.io.IOException;
-import java.security.GeneralSecurityException;
 import java.util.List;
 import java.util.Optional;
-import java.util.function.Function;
 
 public class GcsFileInputPlugin
         implements FileInputPlugin
 {
-    private static final Logger log = Exec.getLogger(GcsFileInputPlugin.class);
-
     @Override
     public ConfigDiff transaction(ConfigSource config,
                                   FileInputPlugin.Control control)
@@ -42,12 +36,12 @@ public class GcsFileInputPlugin
             }
         }
 
-        if (task.getAuthMethod().getString().equals("json_key")) {
+        if (AuthUtils.AuthMethod.json_key.equals(task.getAuthMethod())) {
             if (!task.getJsonKeyfile().isPresent()) {
                 throw new ConfigException("If auth_method is json_key, you have to set json_keyfile");
             }
         }
-        else if (task.getAuthMethod().getString().equals("private_key")) {
+        else if (AuthUtils.AuthMethod.private_key.equals(task.getAuthMethod())) {
             if (!task.getP12Keyfile().isPresent() || !task.getServiceAccountEmail().isPresent()) {
                 throw new ConfigException("If auth_method is private_key, you have to set both service_account_email and p12_keyfile");
             }
@@ -60,11 +54,9 @@ public class GcsFileInputPlugin
             }
         }
 
-        Storage client = GcsFileInput.newGcsClient(task, newGcsAuth(task));
-
         // list files recursively if path_prefix is specified
         if (task.getPathPrefix().isPresent()) {
-            task.setFiles(GcsFileInput.listFiles(task, client));
+            task.setFiles(GcsFileInput.listFiles(task));
         }
         else {
             if (task.getPathFiles().isEmpty()) {
@@ -80,22 +72,6 @@ public class GcsFileInputPlugin
         return resume(task.dump(), task.getFiles().getTaskCount(), control);
     }
 
-    private GcsAuthentication newGcsAuth(PluginTask task)
-    {
-        try {
-            return new GcsAuthentication(
-                    task.getAuthMethod().getString(),
-                    task.getServiceAccountEmail(),
-                    task.getP12Keyfile().map(localFileToPathString()),
-                    task.getJsonKeyfile().map(localFileToPathString()),
-                    task.getApplicationName()
-            );
-        }
-        catch (GeneralSecurityException | IOException ex) {
-            throw new ConfigException(ex);
-        }
-    }
-
     @Override
     public ConfigDiff resume(TaskSource taskSource,
                              int taskCount,
@@ -121,17 +97,6 @@ public class GcsFileInputPlugin
     {
     }
 
-    private Function<LocalFile, String> localFileToPathString()
-    {
-        return new Function<LocalFile, String>()
-        {
-            public String apply(LocalFile file)
-            {
-                return file.getPath().toString();
-            }
-        };
-    }
-
     @Override
     public TransactionalFileInput open(TaskSource taskSource, int taskIndex)
     {

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

@@ -5,13 +5,12 @@ import org.embulk.config.ConfigDefault;
 import org.embulk.config.ConfigInject;
 import org.embulk.config.Task;
 import org.embulk.spi.BufferAllocator;
-import org.embulk.spi.unit.LocalFile;
 
 import java.util.List;
 import java.util.Optional;
 
 public interface PluginTask
-        extends Task, FileList.Task
+        extends Task, AuthUtils.Task, FileList.Task, RetryUtils.Task
 {
     @Config("bucket")
     String getBucket();
@@ -28,44 +27,17 @@ public interface PluginTask
     @ConfigDefault("true")
     boolean getIncremental();
 
-    @Config("auth_method")
-    @ConfigDefault("\"private_key\"")
-    GcsFileInput.AuthMethod getAuthMethod();
-
-    @Config("service_account_email")
-    @ConfigDefault("null")
-    Optional<String> getServiceAccountEmail();
-
     @Config("application_name")
     @ConfigDefault("\"Embulk GCS input plugin\"")
     String getApplicationName();
 
-    // kept for backward compatibility
-    @Config("p12_keyfile_fullpath")
-    @ConfigDefault("null")
-    Optional<String> getP12KeyfileFullpath();
-
-    @Config("p12_keyfile")
-    @ConfigDefault("null")
-    Optional<LocalFile> getP12Keyfile();
-    void setP12Keyfile(Optional<LocalFile> p12Keyfile);
-
-    @Config("json_keyfile")
-    @ConfigDefault("null")
-    Optional<LocalFile> getJsonKeyfile();
-
     @Config("paths")
     @ConfigDefault("[]")
     List<String> getPathFiles();
-    void setPathFiles(List<String> files);
 
     FileList getFiles();
     void setFiles(FileList files);
 
-    @Config("max_connection_retry")
-    @ConfigDefault("10") // 10 times retry to connect GCS server if failed.
-    int getMaxConnectionRetry();
-
     @ConfigInject
     BufferAllocator getBufferAllocator();
 }

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

@@ -0,0 +1,153 @@
+package org.embulk.input.gcs;
+
+import com.google.api.client.auth.oauth2.TokenErrorResponse;
+import com.google.api.client.auth.oauth2.TokenResponseException;
+import com.google.api.client.googleapis.json.GoogleJsonResponseException;
+import com.google.cloud.storage.Blob;
+import com.google.cloud.storage.Storage;
+import org.embulk.config.Config;
+import org.embulk.config.ConfigDefault;
+import org.embulk.spi.Exec;
+import org.embulk.spi.util.RetryExecutor;
+import org.slf4j.Logger;
+
+import java.util.Optional;
+import java.util.function.Predicate;
+
+class RetryUtils
+{
+    interface Task extends org.embulk.config.Task
+    {
+        @Config("max_connection_retry")
+        @ConfigDefault("10") // 10 times retry to connect GCS server if failed.
+        int getMaxConnectionRetry();
+
+        @Config("initial_retry_interval_millis")
+        @ConfigDefault("1000")
+        int getInitialRetryIntervalMillis();
+
+        @Config("maximum_retry_interval_millis")
+        @ConfigDefault("300000")
+        int getMaximumRetryIntervalMillis();
+    }
+
+    private RetryUtils()
+    {
+    }
+
+    private static final Logger LOG = Exec.getLogger(RetryUtils.class);
+
+    /**
+     * A utility predicate to detect status code 4xx of `GoogleJsonResponseException`
+     */
+    private static final Predicate<GoogleJsonResponseException> API_ERROR_NOT_RETRY_4XX = e -> {
+        if (e.getDetails() == null && e.getContent() != null) {
+            LOG.warn("Invalid response was returned : {}", e.getContent());
+            return true;
+        }
+        int statusCode = e.getDetails().getCode();
+        return statusCode / 100 != 4;
+    };
+
+    /**
+     * A utility predicate to detect status code 4xx of `TokenResponseException`
+     * But will retry 400 "Invalid JWS..."
+     */
+    private static final Predicate<TokenResponseException> TOKEN_ERROR_NOT_RETRY_4XX = e -> {
+        Optional<String> errDesc = Optional.ofNullable(e.getDetails()).map(TokenErrorResponse::getErrorDescription);
+        if (errDesc.isPresent()) {
+            // Retry: 400 BadRequest "Invalid JWT..."
+            // Caused by: com.google.api.client.auth.oauth2.TokenResponseException: 400 Bad Request
+            // {
+            //   "error" : "invalid_grant",
+            //   "error_description" : "Invalid JWT: No valid verifier found for issuer."
+            // }
+            if (errDesc.get().contains("Invalid JWT")) {
+                LOG.warn("Invalid response was returned : {}", errDesc.get());
+                return true;
+            }
+        }
+        return e.getStatusCode() / 100 != 4;
+    };
+
+    /**
+     * A default (abstract) retryable impl, which makes use of above 2 predicates
+     * With default behaviors onRetry, etc.
+     *
+     * @param <T>
+     */
+    public abstract static class DefaultRetryable<T> implements RetryExecutor.Retryable<T>
+    {
+        @Override
+        public boolean isRetryableException(Exception exception)
+        {
+            if (exception instanceof GoogleJsonResponseException) {
+                return API_ERROR_NOT_RETRY_4XX.test((GoogleJsonResponseException) exception);
+            }
+            else if (exception instanceof TokenResponseException) {
+                return TOKEN_ERROR_NOT_RETRY_4XX.test((TokenResponseException) exception);
+            }
+            return true;
+        }
+
+        @Override
+        public void onRetry(Exception exception, int retryCount, int retryLimit, int retryWait)
+        {
+            String message = String.format("GCS GET request failed. Retrying %d/%d after %d seconds. Message: %s: %s",
+                    retryCount, retryLimit, retryWait / 1000, exception.getClass(), exception.getMessage());
+            if (retryCount % 3 == 0) {
+                LOG.warn(message, exception);
+            }
+            else {
+                LOG.warn(message);
+            }
+        }
+
+        @Override
+        public void onGiveup(Exception firstException, Exception lastException)
+        {
+        }
+    }
+
+    /**
+     * Return Blob GET op that is ready for {@code withRetry}
+     *
+     * @param client
+     * @param bucket
+     * @param key
+     * @return
+     */
+    static DefaultRetryable<Blob> get(Storage client, String bucket, String key)
+    {
+        return new DefaultRetryable<Blob>()
+        {
+            @Override
+            public Blob call()
+            {
+                return client.get(bucket, key);
+            }
+        };
+    }
+
+    /**
+     * Utility method
+     *
+     * @param task
+     * @param op
+     * @param <T>
+     * @return
+     */
+    static <T> T withRetry(Task task, RetryExecutor.Retryable<T> op)
+    {
+        try {
+            return RetryExecutor.retryExecutor()
+                    .withInitialRetryWait(task.getInitialRetryIntervalMillis())
+                    .withMaxRetryWait(task.getMaximumRetryIntervalMillis())
+                    .withRetryLimit(task.getMaxConnectionRetry())
+                    .runInterruptible(op);
+        }
+        catch (RetryExecutor.RetryGiveupException | InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+    }
+}