embulk-output-bigquery 0.1.0

data/lib/embulk/output/bigquery.rb

@@ -0,0 +1,3 @@
+Embulk::JavaPlugin.register_output(
+  "bigquery", "org.embulk.output.BigqueryOutputPlugin",
+  File.expand_path('../../../../classpath', __FILE__))

data/settings.gradle

@@ -0,0 +1,2 @@
+rootProject.name = 'embulk-output-bigquery'
+

data/src/main/java/org/embulk/output/BigqueryAuthentication.java

@@ -0,0 +1,99 @@
+package org.embulk.output;
+
+import java.io.File;
+import java.io.FileNotFoundException;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.IllegalFormatException;
+import com.google.api.client.auth.oauth2.Credential;
+import com.google.api.client.auth.oauth2.CredentialRefreshListener;
+import com.google.api.client.auth.oauth2.TokenErrorResponse;
+import com.google.api.client.auth.oauth2.TokenResponse;
+import com.google.common.collect.ImmutableList;
+import java.security.GeneralSecurityException;
+
+import org.embulk.spi.Exec;
+import org.slf4j.Logger;
+
+import com.google.api.client.googleapis.auth.oauth2.GoogleCredential;
+import com.google.api.client.googleapis.javanet.GoogleNetHttpTransport;
+import com.google.api.client.http.HttpTransport;
+import com.google.api.client.http.InputStreamContent;
+import com.google.api.client.json.JsonFactory;
+import com.google.api.client.json.jackson2.JacksonFactory;
+import com.google.api.services.storage.Storage;
+import com.google.api.services.storage.StorageScopes;
+import com.google.api.services.bigquery.Bigquery;
+import com.google.api.services.bigquery.BigqueryScopes;
+import com.google.api.services.bigquery.model.ProjectList;
+
+public class BigqueryAuthentication
+{
+
+    private final Logger log = Exec.getLogger(BigqueryAuthentication.class);
+    private final String serviceAccountEmail;
+    private final String p12KeyFilePath;
+    private final String applicationName;
+    private final HttpTransport httpTransport;
+    private final JsonFactory jsonFactory;
+    private final GoogleCredential credentials;
+
+    public BigqueryAuthentication(String serviceAccountEmail, String p12KeyFilePath, String applicationName) throws IOException, GeneralSecurityException
+    {
+        this.serviceAccountEmail = serviceAccountEmail;
+        this.p12KeyFilePath = p12KeyFilePath;
+        this.applicationName = applicationName;
+
+        this.httpTransport = GoogleNetHttpTransport.newTrustedTransport();
+        this.jsonFactory = new JacksonFactory();
+        this.credentials = getCredentialProvider();
+    }
+
+    /**
+     * @see https://developers.google.com/accounts/docs/OAuth2ServiceAccount#authorizingrequests
+     */
+    private GoogleCredential getCredentialProvider() throws IOException, GeneralSecurityException
+    {
+        // @see https://cloud.google.com/compute/docs/api/how-tos/authorization
+        // @see https://developers.google.com/resources/api-libraries/documentation/storage/v1/java/latest/com/google/api/services/storage/STORAGE_SCOPE.html
+        GoogleCredential cred = new GoogleCredential.Builder()
+                .setTransport(httpTransport)
+                .setJsonFactory(jsonFactory)
+                .setServiceAccountId(serviceAccountEmail)
+                .setServiceAccountScopes(
+                        ImmutableList.of(
+                                BigqueryScopes.DEVSTORAGE_READ_WRITE,
+                                BigqueryScopes.BIGQUERY
+                        )
+                )
+                .setServiceAccountPrivateKeyFromP12File(new File(p12KeyFilePath))
+                .build();
+        return cred;
+    }
+
+    public Bigquery getBigqueryClient() throws IOException
+    {
+        Bigquery client = new Bigquery.Builder(httpTransport, jsonFactory, credentials)
+                .setHttpRequestInitializer(credentials)
+                .setApplicationName(applicationName)
+                .build();
+
+        // For throw IOException when authentication is failed.
+        long maxResults = 1;
+        Bigquery.Projects.List req = client.projects().list().setMaxResults(maxResults);
+        ProjectList projectList = req.execute();
+
+        return client;
+    }
+
+    public Storage getGcsClient() throws IOException
+    {
+        Storage client = new Storage.Builder(httpTransport, jsonFactory, credentials)
+                .setApplicationName(applicationName)
+                .build();
+
+        return client;
+    }
+}

data/src/main/java/org/embulk/output/BigqueryGcsWriter.java

@@ -0,0 +1,201 @@
+package org.embulk.output;
+
+import java.io.File;
+import java.io.FileNotFoundException;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Collection;
+import java.util.Iterator;
+import java.util.IllegalFormatException;
+import java.nio.charset.Charset;
+import java.nio.charset.StandardCharsets;
+import com.google.common.base.Optional;
+import com.google.common.collect.ImmutableList;
+//import eu.medsea.mimeutil.MimeType;
+//import eu.medsea.mimeutil.MimeUtil;
+//import eu.medsea.mimeutil.detector.MimeDetector;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.commons.codec.binary.Base64;
+import java.security.GeneralSecurityException;
+
+import org.embulk.spi.Exec;
+import org.slf4j.Logger;
+
+import com.google.api.services.storage.Storage;
+import com.google.api.services.storage.StorageScopes;
+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.client.http.InputStreamContent;
+
+public class BigqueryGcsWriter
+{
+
+    private final Logger log = Exec.getLogger(BigqueryGcsWriter.class);
+    private final String bucket;
+    private final String sourceFormat;
+    private final boolean isFileCompressed;
+    private final boolean deleteFromBucketWhenJobEnd;
+    private Storage storageClient;
+
+    public BigqueryGcsWriter(Builder builder) throws IOException, GeneralSecurityException
+    {
+        this.bucket = builder.bucket;
+        this.sourceFormat = builder.sourceFormat.toUpperCase();
+        this.isFileCompressed = builder.isFileCompressed;
+        this.deleteFromBucketWhenJobEnd = builder.deleteFromBucketWhenJobEnd;
+
+        BigqueryAuthentication auth = new BigqueryAuthentication(builder.serviceAccountEmail, builder.p12KeyFilePath, builder.applicationName);
+        this.storageClient = auth.getGcsClient();
+    }
+
+    public void uploadFile(String localFilePath, String fileName, Optional<String> remotePath) throws IOException
+    {
+        FileInputStream stream = null;
+
+        try {
+            String path;
+            if (remotePath.isPresent()) {
+                path = remotePath.get();
+            } else {
+                path = "";
+            }
+            String gcsPath = getRemotePath(path, fileName);
+            StorageObject objectMetadata = new StorageObject().setName(gcsPath);
+            log.info(String.format("Uploading file [%s] to [gs://%s/%s]", localFilePath, bucket, gcsPath));
+
+            File file = new File(localFilePath);
+            stream = new FileInputStream(file);
+            InputStreamContent content = new InputStreamContent(getContentType(), stream);
+            Storage.Objects.Insert insertObject = storageClient.objects().insert(bucket, objectMetadata, content);
+            insertObject.setDisableGZipContent(true);
+
+            StorageObject response = insertObject.execute();
+            log.info(String.format("Upload completed [%s] to [gs://%s/%s]", localFilePath, bucket, gcsPath));
+        } finally {
+            stream.close();
+        }
+    }
+
+    private String getRemotePath(String remotePath, String fileName)
+    {
+        if (remotePath.isEmpty()) {
+            return fileName;
+        }
+        String[] pathList = StringUtils.split(remotePath, '/');
+        String path = StringUtils.join(pathList) + "/";
+        if (!path.endsWith("/")) {
+            path = path + "/";
+        }
+        return path + fileName;
+    }
+
+    public void deleteFile(String remotePath, String fileName) throws IOException
+    {
+        String path = getRemotePath(remotePath, fileName);
+        storageClient.objects().delete(bucket, path).execute();
+        log.info(String.format("Delete remote file [gs://%s/%s]", bucket, path));
+    }
+
+    public boolean getDeleteFromBucketWhenJobEnd()
+    {
+        return this.deleteFromBucketWhenJobEnd;
+    }
+
+    private String getContentType()
+    {
+        if (isFileCompressed) {
+            return "application/x-gzip";
+        } else {
+            if (sourceFormat.equals("NEWLINE_DELIMITED_JSON)")) {
+                return "application/json";
+            } else {
+                return "text/csv";
+            }
+        }
+    }
+
+    /*
+    private void registerMimeDetector()
+    {
+        String mimeDetector = "eu.medsea.mimeutil.detector.MagicMimeMimeDetector";
+        MimeDetector registeredMimeDetector = MimeUtil.getMimeDetector(mimeDetector);
+        MimeUtil.registerMimeDetector(mimeDetector);
+    }
+
+    public String detectMimeType(File file)
+    {
+        try {
+            Collection<?> mimeTypes = MimeUtil.getMimeTypes(file);
+            if (!mimeTypes.isEmpty()) {
+                Iterator<?> iterator = mimeTypes.iterator();
+                MimeType mimeType = (MimeType) iterator.next();
+                return mimeType.getMediaType() + "/" + mimeType.getSubType();
+            }
+        } catch (Exception ex) {
+        }
+        return "application/octet-stream";
+    }
+    */
+
+    public static class Builder
+    {
+        private final String serviceAccountEmail;
+        private String p12KeyFilePath;
+        private String applicationName;
+        private String bucket;
+        private String sourceFormat;
+        private boolean isFileCompressed;
+        private boolean deleteFromBucketWhenJobEnd;
+        private boolean enableMd5hashCheck;
+
+        public Builder(String serviceAccountEmail)
+        {
+            this.serviceAccountEmail = serviceAccountEmail;
+        }
+
+        public Builder setP12KeyFilePath(String p12KeyFilePath)
+        {
+            this.p12KeyFilePath = p12KeyFilePath;
+            return this;
+        }
+
+        public Builder setApplicationName(String applicationName)
+        {
+            this.applicationName = applicationName;
+            return this;
+        }
+
+        public Builder setBucket(String bucket)
+        {
+            this.bucket = bucket;
+            return this;
+        }
+
+        public Builder setSourceFormat(String sourceFormat)
+        {
+            this.sourceFormat = sourceFormat;
+            return this;
+        }
+
+        public Builder setIsFileCompressed(boolean isFileCompressed)
+        {
+            this.isFileCompressed = isFileCompressed;
+            return this;
+        }
+
+        public Builder setDeleteFromBucketWhenJobEnd(boolean deleteFromBucketWhenJobEnd)
+        {
+            this.deleteFromBucketWhenJobEnd = deleteFromBucketWhenJobEnd;
+            return this;
+        }
+
+        public BigqueryGcsWriter build() throws IOException, GeneralSecurityException
+        {
+            return new BigqueryGcsWriter(this);
+        }
+    }
+}

data/src/main/java/org/embulk/output/BigqueryOutputPlugin.java

@@ -0,0 +1,293 @@
+package org.embulk.output;
+
+import java.io.File;
+import java.io.FileWriter;
+import java.io.FileNotFoundException;
+import java.io.FileOutputStream;
+import java.io.BufferedOutputStream;
+import java.io.IOException;
+import java.util.List;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.concurrent.TimeoutException;
+import com.google.common.base.Optional;
+import com.google.common.base.Throwables;
+import java.security.GeneralSecurityException;
+
+import org.embulk.config.Config;
+import org.embulk.config.ConfigException;
+import org.embulk.config.ConfigDefault;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.ConfigDiff;
+import org.embulk.config.CommitReport;
+import org.embulk.config.Task;
+import org.embulk.config.TaskSource;
+import org.embulk.spi.Buffer;
+import org.embulk.spi.FileOutputPlugin;
+import org.embulk.spi.TransactionalFileOutput;
+import org.embulk.spi.Exec;
+
+import org.slf4j.Logger;
+
+public class BigqueryOutputPlugin
+        implements FileOutputPlugin
+{
+    public interface PluginTask
+            extends Task
+    {
+        @Config("service_account_email")
+        public String getServiceAccountEmail();
+
+        @Config("p12_keyfile_path")
+        public String getP12KeyfilePath();
+
+        @Config("application_name")
+        @ConfigDefault("\"Embulk BigQuery plugin\"")
+        public String getApplicationName();
+
+        @Config("path_prefix")
+        public String getPathPrefix();
+
+        @Config("sequence_format")
+        @ConfigDefault("\".%03d.%02d\"")
+        public String getSequenceFormat();
+
+        @Config("file_ext")
+        public String getFileNameExtension();
+
+        @Config("source_format")
+        @ConfigDefault("\"CSV\"")
+        public String getSourceFormat();
+
+        @Config("is_file_compressed")
+        @ConfigDefault("true")
+        public boolean getIsFileCompressed();
+
+        @Config("field_delimiter")
+        @ConfigDefault("\",\"")
+        public String getFieldDelimiter();
+
+        @Config("max_bad_records")
+        @ConfigDefault("0")
+        public int getMaxBadrecords();
+
+        @Config("delete_from_local_when_upload_end")
+        @ConfigDefault("false")
+        public boolean getDeleteFromLocalWhenUploadEnd();
+
+        @Config("delete_from_bucket_when_job_end")
+        @ConfigDefault("false")
+        public boolean getDeleteFromBucketWhenJobEnd();
+
+        @Config("bucket")
+        public String getBucket();
+
+        @Config("remote_path")
+        @ConfigDefault("null")
+        public Optional<String> getRemotePath();
+
+        @Config("project")
+        public String getProject();
+
+        @Config("dataset")
+        public String getDataset();
+
+        @Config("table")
+        public String getTable();
+
+        @Config("auto_create_table")
+        @ConfigDefault("false")
+        public boolean getAutoCreateTable();
+
+        @Config("schema_path")
+        @ConfigDefault("null")
+        public Optional<String> getSchemaPath();
+
+        @Config("job_status_max_polling_time")
+        @ConfigDefault("3600")
+        public int getJobStatusMaxPollingTime();
+
+        @Config("job_status_polling_interval")
+        @ConfigDefault("10")
+        public int getJobStatusPollingInterval();
+
+        @Config("is_skip_job_result_check")
+        @ConfigDefault("0")
+        public boolean getIsSkipJobResultCheck();
+    }
+
+    private final Logger log = Exec.getLogger(BigqueryOutputPlugin.class);
+    private static BigqueryGcsWriter bigQueryGcsWriter;
+    private static BigqueryWriter bigQueryWriter;
+
+    public ConfigDiff transaction(ConfigSource config, int taskCount,
+                                  FileOutputPlugin.Control control)
+    {
+        final PluginTask task = config.loadConfig(PluginTask.class);
+
+        try {
+            bigQueryGcsWriter = new BigqueryGcsWriter.Builder(task.getServiceAccountEmail())
+                    .setP12KeyFilePath(task.getP12KeyfilePath())
+                    .setApplicationName(task.getApplicationName())
+                    .setBucket(task.getBucket())
+                    .setSourceFormat(task.getSourceFormat())
+                    .setIsFileCompressed(task.getIsFileCompressed())
+                    .setDeleteFromBucketWhenJobEnd(task.getDeleteFromBucketWhenJobEnd())
+                    .build();
+
+            bigQueryWriter = new BigqueryWriter.Builder(task.getServiceAccountEmail())
+                    .setP12KeyFilePath(task.getP12KeyfilePath())
+                    .setApplicationName(task.getApplicationName())
+                    .setProject(task.getProject())
+                    .setDataset(task.getDataset())
+                    .setTable(task.getTable())
+                    .setAutoCreateTable(task.getAutoCreateTable())
+                    .setSchemaPath(task.getSchemaPath())
+                    .setBucket(task.getBucket())
+                    .setSourceFormat(task.getSourceFormat())
+                    .setFieldDelimiter(task.getFieldDelimiter())
+                    .setMaxBadrecords(task.getMaxBadrecords())
+                    .setJobStatusMaxPollingTime(task.getJobStatusMaxPollingTime())
+                    .setJobStatusPollingInterval(task.getJobStatusPollingInterval())
+                    .setIsSkipJobResultCheck(task.getIsSkipJobResultCheck())
+                    .build();
+        } catch (IOException | GeneralSecurityException ex) {
+            log.warn("Google Authentication was failed. Please Check your configurations.");
+            throw new ConfigException(ex);
+        }
+        // non-retryable (non-idempotent) output:
+        return resume(task.dump(), taskCount, control);
+    }
+
+    public ConfigDiff resume(TaskSource taskSource,
+                             int taskCount,
+                             FileOutputPlugin.Control control)
+    {
+        control.run(taskSource);
+
+        try {
+            bigQueryWriter.executeJob();
+            // TODO refactor
+            if (bigQueryGcsWriter.getDeleteFromBucketWhenJobEnd()) {
+                ArrayList<HashMap<String, String>> fileList = bigQueryWriter.getFileList();
+                for (HashMap<String, String> file : fileList) {
+                    bigQueryGcsWriter.deleteFile(file.get("remote_path"), file.get("file_name"));
+                }
+            }
+        } catch (IOException | TimeoutException | BigqueryWriter.JobFailedException ex) {
+            log.warn(ex.getMessage());
+            throw Throwables.propagate(ex);
+        }
+        return Exec.newConfigDiff();
+    }
+
+    @Override
+    public void cleanup(TaskSource taskSource,
+                        int taskCount,
+                        List<CommitReport> successCommitReports)
+    {
+    }
+
+    @Override
+    public TransactionalFileOutput open(TaskSource taskSource, final int taskIndex)
+    {
+        final PluginTask task = taskSource.loadTask(PluginTask.class);
+
+        final String pathPrefix = task.getPathPrefix();
+        final String sequenceFormat = task.getSequenceFormat();
+        final String pathSuffix = task.getFileNameExtension();
+        final Optional<String> remotePath = task.getRemotePath();
+
+        return new TransactionalFileOutput() {
+            private int fileIndex = 0;
+            private BufferedOutputStream output = null;
+            private File file;
+            private String filePath;
+            private String fileName;
+            private long fileSize;
+
+            public void nextFile()
+            {
+                closeFile();
+
+                try {
+                    String suffix = pathSuffix;
+                    if (!suffix.startsWith(".")) {
+                        suffix = "." + suffix;
+                    }
+                    filePath = pathPrefix + String.format(sequenceFormat, taskIndex, fileIndex) + suffix;
+                    file = new File(filePath);
+                    fileName = file.getName();
+                    fileSize = file.length();
+
+                    String parentPath = file.getParent();
+                    File dir = new File(parentPath);
+                    if (!dir.exists()) {
+                        dir.mkdir();
+                    }
+                    log.info(String.format("Writing file [%s]", filePath));
+                    output = new BufferedOutputStream(new FileOutputStream(filePath));
+                } catch (FileNotFoundException ex) {
+                    throw Throwables.propagate(ex);
+                }
+                fileIndex++;
+            }
+
+            private void closeFile()
+            {
+                if (output != null) {
+                    try {
+                        output.close();
+                    } catch (IOException ex) {
+                        throw Throwables.propagate(ex);
+                    }
+                }
+            }
+
+            public void add(Buffer buffer)
+            {
+                try {
+                    output.write(buffer.array(), buffer.offset(), buffer.limit());
+                } catch (IOException ex) {
+                    throw Throwables.propagate(ex);
+                } finally {
+                    buffer.release();
+                }
+            }
+
+            public void finish()
+            {
+                closeFile();
+                if (fileName != null) {
+                    try {
+                        bigQueryGcsWriter.uploadFile(filePath, fileName, remotePath);
+
+                        if (task.getDeleteFromLocalWhenUploadEnd()) {
+                            log.info(String.format("Delete local file [%s]", filePath));
+                            file.delete();
+                        }
+
+                        bigQueryWriter.addTask(remotePath, fileName, fileSize);
+                    } catch (IOException ex) {
+                        throw Throwables.propagate(ex);
+                    }
+                }
+            }
+
+            public void close()
+            {
+                closeFile();
+            }
+
+            public void abort()
+            {
+            }
+
+            public CommitReport commit()
+            {
+                CommitReport report = Exec.newCommitReport();
+                return report;
+            }
+        };
+    }
+}