Handle auth_method & json_keyfile properly

Author: t3t5u

src/main/java/org/embulk/output/bigquery_java/Auth.java

@@ -0,0 +1,45 @@
+package org.embulk.output.bigquery_java;
+
+import com.google.auth.Credentials;
+import com.google.auth.oauth2.ComputeEngineCredentials;
+import com.google.auth.oauth2.GoogleCredentials;
+import com.google.auth.oauth2.ServiceAccountCredentials;
+import com.google.auth.oauth2.UserCredentials;
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import org.embulk.config.ConfigException;
+import org.embulk.output.bigquery_java.config.PluginTask;
+import org.embulk.util.config.units.LocalFile;
+
+public class Auth {
+  private final String authMethod;
+  private final LocalFile jsonKeyFile;
+
+  public Auth(PluginTask task) {
+    authMethod = task.getAuthMethod();
+    jsonKeyFile = task.getJsonKeyfile();
+  }
+
+  public Credentials getCredentials(String... scopes) throws IOException {
+    return getGoogleCredentials().createScoped(scopes);
+  }
+
+  private GoogleCredentials getGoogleCredentials() throws IOException {
+    if ("authorized_user".equalsIgnoreCase(authMethod)) {
+      return UserCredentials.fromStream(getCredentialsStream());
+    } else if ("service_account".equalsIgnoreCase(authMethod)) {
+      return ServiceAccountCredentials.fromStream(getCredentialsStream());
+    } else if ("compute_engine".equalsIgnoreCase(authMethod)) {
+      return ComputeEngineCredentials.create();
+    } else if ("application_default".equalsIgnoreCase(authMethod)) {
+      return GoogleCredentials.getApplicationDefault();
+    } else {
+      throw new ConfigException("Unknown auth method: " + authMethod);
+    }
+  }
+
+  private InputStream getCredentialsStream() {
+    return new ByteArrayInputStream(jsonKeyFile.getContent());
+  }
+}

src/main/java/org/embulk/output/bigquery_java/BigqueryClient.java

@@ -1,6 +1,6 @@
 package org.embulk.output.bigquery_java;
 
-import com.google.auth.oauth2.ServiceAccountCredentials;
+import com.google.api.services.bigquery.BigqueryScopes;
 import com.google.cloud.bigquery.BigQuery;
 import com.google.cloud.bigquery.BigQueryException;
 import com.google.cloud.bigquery.BigQueryOptions;
@@ -27,7 +27,6 @@
 import com.google.cloud.bigquery.TableResult;
 import com.google.cloud.bigquery.TimePartitioning;
 import com.google.cloud.bigquery.WriteChannelConfiguration;
-import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.channels.Channels;
@@ -88,15 +87,15 @@ public BigqueryClient(PluginTask task, Schema schema) {
     }
     columnOptions = task.getColumnOptions().orElse(Collections.emptyList());
     try {
-      bigquery = getClientWithJsonKey(task.getJsonKeyfile());
+      bigquery = getBigQueryService();
     } catch (IOException e) {
       throw new RuntimeException(e);
     }
   }
 
-  private static BigQuery getClientWithJsonKey(String key) throws IOException {
+  private BigQuery getBigQueryService() throws IOException {
     return BigQueryOptions.newBuilder()
-        .setCredentials(ServiceAccountCredentials.fromStream(new FileInputStream(key)))
+        .setCredentials(new Auth(task).getCredentials(BigqueryScopes.BIGQUERY))
         .build()
         .getService();
   }

src/main/java/org/embulk/output/bigquery_java/config/PluginTask.java

@@ -5,6 +5,7 @@
 import org.embulk.util.config.Config;
 import org.embulk.util.config.ConfigDefault;
 import org.embulk.util.config.Task;
+import org.embulk.util.config.units.LocalFile;
 
 public interface PluginTask extends Task {
 
@@ -20,7 +21,7 @@ public interface PluginTask extends Task {
   String getAuthMethod();
 
   @Config("json_keyfile")
-  String getJsonKeyfile();
+  LocalFile getJsonKeyfile();
 
   @Config("dataset")
   String getDataset();

src/test/java/org/embulk/output/bigquery_java/TestBigqueryJavaOutputPlugin.java

@@ -26,6 +26,7 @@
 import org.embulk.util.config.ConfigMapper;
 import org.embulk.util.config.ConfigMapperFactory;
 import org.embulk.util.config.Task;
+import org.embulk.util.config.units.LocalFile;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
@@ -85,7 +86,7 @@ public void testWithTimePartitioning() {
 
   public interface TestTask extends Task {
     @Config("json_keyfile")
-    String getJsonKeyfile();
+    LocalFile getJsonKeyfile();
 
     @Config("dataset")
     String getDataset();

src/test/java/org/embulk/output/bigquery_java/config/TestBigqueryTaskBuilder.java

@@ -43,7 +43,7 @@ public void setAbortOnError_DefaultMaxBadRecord_True() {
                     "type: bigquery_java",
                     "mode: replace",
                     "auth_method: service_account",
-                    "json_keyfile: json_key.json",
+                    "json_keyfile: { content: \"\" }",
                     "dataset: dataset",
                     "table: table",
                     "source_format: NEWLINE_DELIMITED_JSON",
@@ -71,7 +71,7 @@ public void setFileExt_JSONL_GZIP_JSONL_GZ() {
                     "type: bigquery_java",
                     "mode: replace",
                     "auth_method: service_account",
-                    "json_keyfile: json_key.json",
+                    "json_keyfile: { content: \"\" }",
                     "dataset: dataset",
                     "table: table",
                     "source_format: NEWLINE_DELIMITED_JSON",
@@ -99,7 +99,7 @@ public void clustering() {
                     "type: bigquery_java",
                     "mode: replace",
                     "auth_method: service_account",
-                    "json_keyfile: json_key.json",
+                    "json_keyfile: { content: \"\" }",
                     "dataset: dataset",
                     "table: table",
                     "source_format: NEWLINE_DELIMITED_JSON",

src/test/java/org/embulk/output/bigquery_java/visitor/TestJsonColumnVisitor.java

@@ -20,6 +20,7 @@
 import org.embulk.spi.type.Types;
 import org.embulk.util.config.ConfigMapper;
 import org.embulk.util.config.ConfigMapperFactory;
+import org.embulk.util.config.units.LocalFile;
 import org.junit.Test;
 import org.msgpack.value.Value;
 import org.msgpack.value.impl.ImmutableArrayValueImpl;
@@ -130,7 +131,7 @@ private Object visitColumn(
       throws JsonProcessingException {
     ConfigSource configSource = CONFIG_MAPPER_FACTORY.newConfigSource();
     configSource.set("mode", "replace");
-    configSource.set("json_keyfile", "test");
+    configSource.set("json_keyfile", LocalFile.ofContent(""));
     configSource.set("dataset", "test");
     configSource.set("table", "test");
     configSource.set("source_format", "NEWLINE_DELIMITED_JSON");

src/test/resources/java/org/embulk/output/bigquery_java/base.yml

@@ -1,7 +1,7 @@
 type: bigquery_java
 mode: replace
 auth_method: service_account
-json_keyfile: json_key.json
+json_keyfile: { content: "" }
 dataset: dataset
 table: table
 source_format: NEWLINE_DELIMITED_JSON

src/test/resources/java/org/embulk/output/bigquery_java/time_partitioning.yml

@@ -1,7 +1,7 @@
 type: bigquery_java
 mode: replace
 auth_method: service_account
-json_keyfile: json_key.json
+json_keyfile: { content: "" }
 dataset: dataset
 table: table
 source_format: NEWLINE_DELIMITED_JSON