embulk-output-redshift 0.5.0 → 0.5.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: fae200c792daa5799eea798de528bfe89716cad7
-  data.tar.gz: eaf7d9294216da01c956b3498165fbb876f80a11
+  metadata.gz: c8c2e9f95d662d1860ccba4b210c6dae2d1cc44b
+  data.tar.gz: cd0afa3f4352bc9732e27d1fdd132398114356c6
 SHA512:
-  metadata.gz: c47c00911b4bf7a34a38994afef62e21a3695c8ea990e4c1c45b2db2102eb890b3480e9f62677c60a622a47b796e67816adfeddfb8dd601b7586e60e9fb1c8e9
-  data.tar.gz: 21d1dcec368cc4be9cf3e29e001605bf06e7bd727873d8685c409faa257599c6e041cc77aebf161466be8f5bbb3ee0bf054db860bf5fbcd17046b42ff72a4891
+  metadata.gz: 15a7eb5c206c87822f163dd105af21444cf3d925f01e01e0eb7e4ab0e499646b814e9f9b377b49df677b79cb7a559ac97674f7179f699f2dc3cd96523e99216c
+  data.tar.gz: dd060ac7a977ee44b20e5927c3e121d2b930e81b18282b0407a658921db8c7f9fb392fdb277c2078b9dbf08f13de3d987337cb2064c7637205726b82035f85f9

data/README.md

@@ -5,8 +5,8 @@ Redshift output plugins for Embulk loads records to Redshift.
 ## Overview
 
 * **Plugin type**: output
-* **Load all or nothing**: depnds on the mode. see bellow.
-* **Resume supported**: depnds on the mode. see bellow.
+* **Load all or nothing**: depnds on the mode. see below.
+* **Resume supported**: depnds on the mode. see below.
 
 ## Configuration
 
@@ -19,11 +19,11 @@ Redshift output plugins for Embulk loads records to Redshift.
 - **table**: destination table name (string, required)
 - **access_key_id**: access key id for AWS
 - **secret_access_key**: secret access key for AWS
-- **iam_user_name**: IAM user name for uploading temporary files to S3. The user should have permissions of `s3:GetObject`, `s3:PutObject`, `s3:ListBucket` and `sts:GetFederationToken`.
+- **iam_user_name**: IAM user name for uploading temporary files to S3. The user should have permissions of `s3:GetObject`, `s3:PutObject`, `s3:DeleteObject`, , `s3:ListBucket` and `sts:GetFederationToken`. (string, default: "", but we strongly recommend that you use IAM user for security reasons. see below.)
 - **s3_bucket**: S3 bucket name for temporary files
 - **s3_key_prefix**: S3 key prefix for temporary files (string, default:"")
 - **options**: extra connection properties (hash, default: {})
-- **mode**: "replace" or "insert" (string, required)
+- **mode**: "insert", "insert_direct", "truncate_insert", or "replace". See below. (string, required)
 - **batch_size**: size of a single batch insert (integer, default: 16777216)
 - **default_timezone**: If input column type (embulk type) is timestamp, this plugin needs to format the timestamp into a SQL string. This default_timezone option is used to control the timezone. You can overwrite timezone for each columns using column_options option. (string, default: `UTC`)
 - **column_options**: advanced: a key-value pairs where key is a column name and value is options for the column.
@@ -32,6 +32,7 @@ Redshift output plugins for Embulk loads records to Redshift.
   - **timestamp_format**: If input column type (embulk type) is timestamp and value_type is `string` or `nstring`, this plugin needs to format the timestamp value into a string. This timestamp_format option is used to control the format of the timestamp. (string, default: `%Y-%m-%d %H:%M:%S.%6N`)
   - **timezone**: If input column type (embulk type) is timestamp, this plugin needs to format the timestamp value into a SQL string. In this cases, this timezone option is used to control the timezone. (string, value of default_timezone option is used by default)
 
+
 ### Modes
 
 * **insert**:
@@ -98,3 +99,8 @@ out:
 ```
 $ ./gradlew gem
 ```
+
+### Security
+This plugin requires AWS access credentials so that it may write temporary files to S3. There are two security options, Standard and Federated.
+To use Standard security, give **aws_key_id** and **secret_access_key**. To use Federated mode, also give the **iam_user_name** field.
+Federated mode really means temporary credentials, so that a man-in-the-middle attack will see AWS credentials that are only valid for 1 calendar day after the transaction.

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

@@ -53,6 +53,7 @@ public class RedshiftOutputPlugin
         public String getSecretAccessKey();
 
         @Config("iam_user_name")
+        @ConfigDefault("\"\"")
         public String getIamUserName();
 
         @Config("s3_bucket")

data/src/main/java/org/embulk/output/redshift/RedshiftCopyBatchInsert.java

@@ -1,15 +1,27 @@
 package org.embulk.output.redshift;
 
-import java.util.zip.GZIPOutputStream;
-import java.util.concurrent.Callable;
-import java.util.UUID;
+import java.io.BufferedWriter;
 import java.io.File;
-import java.io.IOException;
 import java.io.FileOutputStream;
+import java.io.IOException;
 import java.io.OutputStreamWriter;
-import java.io.Closeable;
-import java.io.BufferedWriter;
 import java.sql.SQLException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.UUID;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.zip.GZIPOutputStream;
+
+import org.embulk.output.jdbc.JdbcSchema;
+import org.embulk.output.postgresql.AbstractPostgreSQLCopyBatchInsert;
+import org.embulk.spi.Exec;
+import org.slf4j.Logger;
+
 import com.amazonaws.auth.AWSCredentialsProvider;
 import com.amazonaws.auth.BasicSessionCredentials;
 import com.amazonaws.auth.policy.Policy;
@@ -19,13 +31,9 @@ import com.amazonaws.auth.policy.Statement.Effect;
 import com.amazonaws.auth.policy.actions.S3Actions;
 import com.amazonaws.services.s3.AmazonS3Client;
 import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClient;
+import com.amazonaws.services.securitytoken.model.Credentials;
 import com.amazonaws.services.securitytoken.model.GetFederationTokenRequest;
 import com.amazonaws.services.securitytoken.model.GetFederationTokenResult;
-import com.amazonaws.services.securitytoken.model.Credentials;
-import org.slf4j.Logger;
-import org.embulk.spi.Exec;
-import org.embulk.output.jdbc.JdbcSchema;
-import org.embulk.output.postgresql.AbstractPostgreSQLCopyBatchInsert;
 
 public class RedshiftCopyBatchInsert
         extends AbstractPostgreSQLCopyBatchInsert
@@ -35,13 +43,16 @@ public class RedshiftCopyBatchInsert
     private final String s3BucketName;
     private final String s3KeyPrefix;
     private final String iamReaderUserName;
+    private final AWSCredentialsProvider credentialsProvider;
     private final AmazonS3Client s3;
     private final AWSSecurityTokenServiceClient sts;
+    private final ExecutorService executorService;
 
     private RedshiftOutputConnection connection = null;
     private String copySqlBeforeFrom = null;
     private long totalRows;
     private int fileCount;
+    private List<Future<Void>> uploadAndCopyFutures;
 
     public static final String COPY_AFTER_FROM = "GZIP DELIMITER '\\t' NULL '\\\\N' ESCAPE TRUNCATECOLUMNS ACCEPTINVCHARS STATUPDATE OFF COMPUPDATE OFF";
 
@@ -58,8 +69,12 @@ public class RedshiftCopyBatchInsert
             this.s3KeyPrefix = s3KeyPrefix + "/";
         }
         this.iamReaderUserName = iamReaderUserName;
+        this.credentialsProvider = credentialsProvider;
         this.s3 = new AmazonS3Client(credentialsProvider);  // TODO options
         this.sts = new AWSSecurityTokenServiceClient(credentialsProvider);  // options
+
+        this.executorService = Executors.newCachedThreadPool();
+        this.uploadAndCopyFutures = new ArrayList<Future<Void>>();
     }
 
     @Override
@@ -86,28 +101,51 @@ public class RedshiftCopyBatchInsert
     {
         File file = closeCurrentFile();  // flush buffered data in writer
 
-        // TODO multi-threading
-        new UploadAndCopyTask(file, batchRows, s3KeyPrefix + UUID.randomUUID().toString()).call();
-        new DeleteFileFinalizer(file).close();
+        String s3KeyName = s3KeyPrefix + UUID.randomUUID().toString();
+        UploadTask uploadTask = new UploadTask(file, batchRows, s3KeyName);
+        Future<Void> uploadFuture = executorService.submit(uploadTask);
+        uploadAndCopyFutures.add(uploadFuture);
+
+        CopyTask copyTask = new CopyTask(uploadFuture, s3KeyName);
+        uploadAndCopyFutures.add(executorService.submit(copyTask));
 
         fileCount++;
         totalRows += batchRows;
         batchRows = 0;
 
         openNewFile();
-        file.delete();
     }
 
     @Override
     public void finish() throws IOException, SQLException
     {
         super.finish();
+
+        for (Future<Void> uploadAndCopyFuture : uploadAndCopyFutures) {
+            try {
+                uploadAndCopyFuture.get();
+
+            } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            } catch (ExecutionException e) {
+                if (e.getCause() instanceof SQLException) {
+                    throw (SQLException)e.getCause();
+                }
+                throw new RuntimeException(e);
+            }
+        }
+
         logger.info("Loaded {} files.", fileCount);
     }
 
     @Override
     public void close() throws IOException, SQLException
     {
+        executorService.shutdownNow();
+        try {
+            executorService.awaitTermination(60, TimeUnit.SECONDS);
+        } catch (InterruptedException e) {}
+
         s3.shutdown();
         closeCurrentFile().delete();
         if (connection != null) {
@@ -127,60 +165,97 @@ public class RedshiftCopyBatchInsert
                         .withActions(S3Actions.GetObject)
                         .withResources(new Resource("arn:aws:s3:::"+s3BucketName+"/"+s3KeyName))  // TODO encode file name using percent encoding
                     );
-        GetFederationTokenRequest req = new GetFederationTokenRequest();
-        req.setDurationSeconds(86400);  // 3600 - 129600
-        req.setName(iamReaderUserName);
-        req.setPolicy(policy.toJson());
-
-        GetFederationTokenResult res = sts.getFederationToken(req);
-        Credentials c = res.getCredentials();
-
-        return new BasicSessionCredentials(
-                c.getAccessKeyId(),
-                c.getSecretAccessKey(),
-                c.getSessionToken());
+        if (iamReaderUserName != null && iamReaderUserName.length() > 0) {
+            GetFederationTokenRequest req = new GetFederationTokenRequest();
+            req.setDurationSeconds(86400);  // 3600 - 129600
+            req.setName(iamReaderUserName);
+            req.setPolicy(policy.toJson());
+
+            GetFederationTokenResult res = sts.getFederationToken(req);
+            Credentials c = res.getCredentials();
+
+            return new BasicSessionCredentials(
+                    c.getAccessKeyId(),
+                    c.getSecretAccessKey(),
+                    c.getSessionToken());
+        } else {
+            return new BasicSessionCredentials(credentialsProvider.getCredentials().getAWSAccessKeyId(),
+                    credentialsProvider.getCredentials().getAWSSecretKey(), null);
+        }
     }
 
-    private class UploadAndCopyTask implements Callable<Void>
+    private class UploadTask implements Callable<Void>
     {
         private final File file;
         private final int batchRows;
         private final String s3KeyName;
 
-        public UploadAndCopyTask(File file, int batchRows, String s3KeyName)
+        public UploadTask(File file, int batchRows, String s3KeyName)
         {
             this.file = file;
             this.batchRows = batchRows;
             this.s3KeyName = s3KeyName;
         }
 
-        public Void call() throws SQLException {
+        public Void call() {
             logger.info(String.format("Uploading file id %s to S3 (%,d bytes %,d rows)",
                         s3KeyName, file.length(), batchRows));
-            s3.putObject(s3BucketName, s3KeyName, file);
 
-            RedshiftOutputConnection con = connector.connect(true);
             try {
-                logger.info("Running COPY from file {}", s3KeyName);
-
-                // create temporary credential right before COPY operation because
-                // it has timeout.
-                // TODO skip this step if iamReaderUserName is not set
-                BasicSessionCredentials creds = generateReaderSessionCredentials(s3KeyName);
-
                 long startTime = System.currentTimeMillis();
-                con.runCopy(buildCopySQL(creds));
+                s3.putObject(s3BucketName, s3KeyName, file);
                 double seconds = (System.currentTimeMillis() - startTime) / 1000.0;
 
-                logger.info(String.format("Loaded file %s (%.2f seconds for COPY)", s3KeyName, seconds));
+                logger.info(String.format("Uploaded file %s (%.2f seconds)", s3KeyName, seconds));
+            } finally {
+                file.delete();
+            }
+
+            return null;
+        }
+    }
 
+    private class CopyTask implements Callable<Void>
+    {
+        private final Future<Void> uploadFuture;
+        private final String s3KeyName;
+
+        public CopyTask(Future<Void> uploadFuture, String s3KeyName)
+        {
+            this.uploadFuture = uploadFuture;
+            this.s3KeyName = s3KeyName;
+        }
+
+        public Void call() throws SQLException, InterruptedException, ExecutionException {
+            try {
+                uploadFuture.get();
+
+                RedshiftOutputConnection con = connector.connect(true);
+                try {
+                    logger.info("Running COPY from file {}", s3KeyName);
+
+                    // create temporary credential right before COPY operation because
+                    // it has timeout.
+                    BasicSessionCredentials creds = generateReaderSessionCredentials(s3KeyName);
+
+                    long startTime = System.currentTimeMillis();
+                    con.runCopy(buildCopySQL(creds));
+                    double seconds = (System.currentTimeMillis() - startTime) / 1000.0;
+
+                    logger.info(String.format("Loaded file %s (%.2f seconds for COPY)", s3KeyName, seconds));
+
+                } finally {
+                    con.close();
+                }
             } finally {
-                con.close();
+                s3.deleteObject(s3BucketName, s3KeyName);
             }
 
             return null;
         }
 
+
+
         private String buildCopySQL(BasicSessionCredentials creds)
         {
             StringBuilder sb = new StringBuilder();
@@ -194,25 +269,13 @@ public class RedshiftCopyBatchInsert
             sb.append(creds.getAWSAccessKeyId());
             sb.append(";aws_secret_access_key=");
             sb.append(creds.getAWSSecretKey());
-            sb.append(";token=");
-            sb.append(creds.getSessionToken());
+            if (creds.getSessionToken() != null) {
+                sb.append(";token=");
+                sb.append(creds.getSessionToken());
+            }
             sb.append("' ");
             sb.append(COPY_AFTER_FROM);
             return sb.toString();
         }
     }
-
-    private static class DeleteFileFinalizer implements Closeable
-    {
-        private File file;
-
-        public DeleteFileFinalizer(File file) {
-            this.file = file;
-        }
-
-        @Override
-        public void close() throws IOException {
-            file.delete();
-        }
-    }
 }

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: embulk-output-redshift
 version: !ruby/object:Gem::Version
-  version: 0.5.0
+  version: 0.5.1
 platform: ruby
 authors:
 - Sadayuki Furuhashi
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2016-01-15 00:00:00.000000000 Z
+date: 2016-03-29 00:00:00.000000000 Z
 dependencies: []
 description: Inserts or updates records to a table.
 email:
@@ -30,9 +30,9 @@ files:
 - classpath/aws-java-sdk-sts-1.10.33.jar
 - classpath/commons-codec-1.6.jar
 - classpath/commons-logging-1.1.3.jar
-- classpath/embulk-output-jdbc-0.5.0.jar
-- classpath/embulk-output-postgresql-0.5.0.jar
-- classpath/embulk-output-redshift-0.5.0.jar
+- classpath/embulk-output-jdbc-0.5.1.jar
+- classpath/embulk-output-postgresql-0.5.1.jar
+- classpath/embulk-output-redshift-0.5.1.jar
 - classpath/httpclient-4.3.6.jar
 - classpath/httpcore-4.3.3.jar
 - classpath/postgresql-9.4-1205-jdbc41.jar