embulk-output-redshift 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA1:
+  metadata.gz: 2824d8d4c60e62b49127ff8759a1c5e50b76e9bc
+  data.tar.gz: 6eccc67a5ac7a8971f5dd4e0a6fabc58d2ade100
+SHA512:
+  metadata.gz: 0a6b84c2d7e0f6d9aca52604738a99e1d0ec6c60033f45f063f807766d93c05a704fc2bf945ade880352ba623d9b4f046de1b9f3b3ec50b32bef4c1dad573406
+  data.tar.gz: 326d1ed8ab544b6740aefa9cd1d7c4459976acfdd57c78b08b411541c36e6b1dbfe56d92a5f923936864df7564e5d7bb19cb4020307c74071a4e01b7a5bb1026

data/README.md

@@ -0,0 +1,47 @@
+# Redshift output plugins for Embulk
+
+Redshift output plugins for Embulk loads records to Redshift.
+
+## Overview
+
+* **Plugin type**: output
+* **Load all or nothing**: depnds on the mode:
+  * **insert**: no
+  * **replace**: yes
+* **Resume supported**: no
+
+## Configuration
+
+- **host**: database host name (string, required)
+- **port**: database port number (integer, default: 5439)
+- **user**: database login user name (string, required)
+- **password**: database login password (string, default: "")
+- **database**: destination database name (string, required)
+- **schema**: destination name (string, default: "public")
+- **table**: destination name (string, required)
+- **mode**: "replace" or "insert" (string, required)
+- **batch_size**: size of a single batch insert (integer, default: 16777216)
+- **options**: extra connection properties (hash, default: {})
+
+### Example
+
+```yaml
+out:
+  type: redshift
+  host: myinstance.us-west-2.redshift.amazonaws.com
+  user: pg
+  password: ""
+  database: my_database
+  table: my_table
+  access_key_id: ABCXYZ123ABCXYZ123
+  secret_access_key: AbCxYz123aBcXyZ123
+  s3_bucket: my-redshift-transfer-bucket
+  iam_user_name: my-s3-read-only
+  mode: insert
+```
+
+### Build
+
+```
+$ ./gradlew gem
+```

data/build.gradle

@@ -0,0 +1,9 @@
+dependencies {
+    compile project(':embulk-output-jdbc')
+    compile project(':embulk-output-postgresql')
+
+    compile "com.amazonaws:aws-java-sdk-s3:1.9.17"
+    compile "com.amazonaws:aws-java-sdk-sts:1.9.17"
+
+    testCompile project(':embulk-output-jdbc').sourceSets.test.output
+}

data/lib/embulk/output/redshift.rb

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

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

@@ -0,0 +1,92 @@
+package org.embulk.output;
+
+import java.util.Properties;
+import java.io.IOException;
+import java.sql.SQLException;
+import org.slf4j.Logger;
+import com.amazonaws.auth.AWSCredentials;
+import com.amazonaws.auth.BasicAWSCredentials;
+import org.embulk.spi.Exec;
+import org.embulk.config.Config;
+import org.embulk.output.jdbc.AbstractJdbcOutputPlugin;
+import org.embulk.output.jdbc.BatchInsert;
+import org.embulk.output.redshift.RedshiftOutputConnector;
+import org.embulk.output.redshift.RedshiftCopyBatchInsert;
+
+public class RedshiftOutputPlugin
+        extends AbstractJdbcOutputPlugin
+{
+    private static final String DEFAULT_SCHEMA = "public";
+    private static final int DEFAULT_PORT = 5439;
+
+    private final Logger logger = Exec.getLogger(RedshiftOutputPlugin.class);
+
+    public interface RedshiftPluginTask extends PluginTask
+    {
+        @Config("access_key_id")
+        public String getAccessKeyId();
+
+        @Config("secret_access_key")
+        public String getSecretAccessKey();
+
+        @Config("iam_user_name")
+        public String getIamUserName();
+
+        @Config("s3_bucket")
+        public String getS3Bucket();
+    }
+
+    @Override
+    protected Class<? extends PluginTask> getTaskClass()
+    {
+        return RedshiftPluginTask.class;
+    }
+
+    @Override
+    protected RedshiftOutputConnector getConnector(PluginTask task, boolean retryableMetadataOperation)
+    {
+        String url = String.format("jdbc:postgresql://%s:%d/%s",
+                task.getHost(), task.getPort().or(DEFAULT_PORT), task.getDatabase());
+
+        Properties props = new Properties();
+        props.setProperty("user", task.getUser());
+        props.setProperty("password", task.getPassword());
+        props.setProperty("loginTimeout",   "300"); // seconds
+        props.setProperty("socketTimeout", "1800"); // seconds
+
+        // Enable keepalive based on tcp_keepalive_time, tcp_keepalive_intvl and tcp_keepalive_probes kernel parameters.
+        // Socket options TCP_KEEPCNT, TCP_KEEPIDLE, and TCP_KEEPINTVL are not configurable.
+        props.setProperty("tcpKeepAlive", "true");
+
+        // TODO
+        //switch task.getSssl() {
+        //when "disable":
+        //    break;
+        //when "enable":
+        //    props.setProperty("sslfactory", "org.postgresql.ssl.NonValidatingFactory");  // disable server-side validation
+        //when "verify":
+        //    props.setProperty("ssl", "true");
+        //    break;
+        //}
+
+        if (!retryableMetadataOperation) {
+            // non-retryable batch operation uses longer timeout
+            props.setProperty("loginTimeout",    "300");  // seconds
+            props.setProperty("socketTimeout", "28800");  // seconds
+        }
+
+        props.putAll(task.getOptions());
+
+        return new RedshiftOutputConnector(url, props, task.getSchema().or(DEFAULT_SCHEMA));
+    }
+
+    @Override
+    protected BatchInsert newBatchInsert(PluginTask task) throws IOException, SQLException
+    {
+        RedshiftPluginTask rt = (RedshiftPluginTask) task;
+        AWSCredentials creds = new BasicAWSCredentials(
+                rt.getAccessKeyId(), rt.getSecretAccessKey());
+        return new RedshiftCopyBatchInsert(getConnector(task, true),
+                creds, rt.getS3Bucket(), rt.getIamUserName());
+    }
+}

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

@@ -0,0 +1,216 @@
+package org.embulk.output.redshift;
+
+import java.util.zip.GZIPOutputStream;
+import java.util.concurrent.Callable;
+import java.util.UUID;
+import java.io.File;
+import java.io.IOException;
+import java.io.FileOutputStream;
+import java.io.OutputStreamWriter;
+import java.io.Closeable;
+import java.io.Writer;
+import java.io.BufferedWriter;
+import java.sql.Connection;
+import java.sql.SQLException;
+import com.amazonaws.auth.AWSCredentials;
+import com.amazonaws.auth.BasicSessionCredentials;
+import com.amazonaws.auth.policy.Policy;
+import com.amazonaws.auth.policy.Resource;
+import com.amazonaws.auth.policy.Statement;
+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.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
+{
+    private final Logger logger = Exec.getLogger(RedshiftCopyBatchInsert.class);
+    private final RedshiftOutputConnector connector;
+    private final AWSCredentials awsCredentials;
+    private final String s3BucketName;
+    private final String iamReaderUserName;
+    private final AmazonS3Client s3;
+    private final AWSSecurityTokenServiceClient sts;
+
+    private RedshiftOutputConnection connection = null;
+    private String copySqlBeforeFrom = null;
+    private long totalRows;
+    private int fileCount;
+
+    public static final String COPY_AFTER_FROM = "GZIP DELIMITER '\\t' NULL '\\N' ESCAPE TRUNCATECOLUMNS ACCEPTINVCHARS STATUPDATE OFF COMPUPDATE OFF";
+
+    public RedshiftCopyBatchInsert(RedshiftOutputConnector connector,
+            AWSCredentials awsCredentials, String s3BucketName,
+            String iamReaderUserName) throws IOException, SQLException
+    {
+        super();
+        this.connector = connector;
+        this.awsCredentials = awsCredentials;
+        this.s3BucketName = s3BucketName;
+        this.iamReaderUserName = iamReaderUserName;
+        this.s3 = new AmazonS3Client(awsCredentials);  // TODO options
+        this.sts = new AWSSecurityTokenServiceClient(awsCredentials);  // options
+    }
+
+    @Override
+    public void prepare(String loadTable, JdbcSchema insertSchema) throws SQLException
+    {
+        this.connection = connector.connect(true);
+        this.copySqlBeforeFrom = connection.buildCopySQLBeforeFrom(loadTable, insertSchema);
+        logger.info("Copy SQL: "+copySqlBeforeFrom+" ? "+COPY_AFTER_FROM);
+    }
+
+    @Override
+    protected BufferedWriter openWriter(File newFile) throws IOException
+    {
+        // Redshift supports gzip
+        return new BufferedWriter(
+                new OutputStreamWriter(
+                    new GZIPOutputStream(new FileOutputStream(newFile)),
+                    FILE_CHARSET)
+                );
+    }
+
+    @Override
+    public void flush() throws IOException, SQLException
+    {
+        File file = closeCurrentFile();  // flush buffered data in writer
+
+        // TODO multi-threading
+        new UploadAndCopyTask(file, batchRows, UUID.randomUUID().toString()).call();
+        new DeleteFileFinalizer(file).close();
+
+        fileCount++;
+        totalRows += batchRows;
+        batchRows = 0;
+
+        openNewFile();
+        file.delete();
+    }
+
+    @Override
+    public void finish() throws IOException, SQLException
+    {
+        super.finish();
+        logger.info("Loaded {} files.", fileCount);
+    }
+
+    @Override
+    public void close() throws IOException, SQLException
+    {
+        s3.shutdown();
+        closeCurrentFile().delete();
+        if (connection != null) {
+            connection.close();
+            connection = null;
+        }
+    }
+
+    private BasicSessionCredentials generateReaderSessionCredentials(String s3KeyName)
+    {
+        Policy policy = new Policy()
+            .withStatements(
+                    new Statement(Effect.Allow)
+                        .withActions(S3Actions.ListObjects)
+                        .withResources(new Resource("arn:aws:s3:::"+s3BucketName)),
+                    new Statement(Effect.Allow)
+                        .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());
+    }
+
+    private class UploadAndCopyTask implements Callable<Void>
+    {
+        private final File file;
+        private final int batchRows;
+        private final String s3KeyName;
+
+        public UploadAndCopyTask(File file, int batchRows, String s3KeyName)
+        {
+            this.file = file;
+            this.batchRows = batchRows;
+            this.s3KeyName = s3KeyName;
+        }
+
+        public Void call() throws SQLException {
+            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));
+                double seconds = (System.currentTimeMillis() - startTime) / 1000.0;
+
+                logger.info(String.format("Loaded file %s (%.2f seconds for COPY)", s3KeyName, seconds));
+
+            } finally {
+                con.close();
+            }
+
+            return null;
+        }
+
+        private String buildCopySQL(BasicSessionCredentials creds)
+        {
+            StringBuilder sb = new StringBuilder();
+            sb.append(copySqlBeforeFrom);
+            sb.append(" FROM 's3://");
+            sb.append(s3BucketName);
+            sb.append("/");
+            sb.append(s3KeyName);
+            sb.append("' CREDENTIALS '");
+            sb.append("aws_access_key_id=");
+            sb.append(creds.getAWSAccessKeyId());
+            sb.append(";aws_secret_access_key=");
+            sb.append(creds.getAWSSecretKey());
+            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();
+        }
+    }
+}

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

@@ -0,0 +1,123 @@
+package org.embulk.output.redshift;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.sql.Statement;
+import org.slf4j.Logger;
+import org.embulk.spi.Exec;
+import org.embulk.output.jdbc.JdbcOutputConnection;
+import org.embulk.output.jdbc.JdbcColumn;
+import org.embulk.output.jdbc.JdbcSchema;
+
+public class RedshiftOutputConnection
+        extends JdbcOutputConnection
+{
+    private final Logger logger = Exec.getLogger(RedshiftOutputConnection.class);
+
+    public RedshiftOutputConnection(Connection connection, String schemaName, boolean autoCommit)
+            throws SQLException
+    {
+        super(connection, schemaName);
+        connection.setAutoCommit(autoCommit);
+    }
+
+    // Redshift does not support DROP TABLE IF EXISTS.
+    // Here runs DROP TABLE and ignores errors.
+    @Override
+    public void dropTableIfExists(String tableName) throws SQLException
+    {
+        Statement stmt = connection.createStatement();
+        try {
+            String sql = String.format("DROP TABLE IF EXISTS %s", quoteIdentifierString(tableName));
+            executeUpdate(stmt, sql);
+            connection.commit();
+        } catch (SQLException ex) {
+            // ignore errors.
+            // TODO here should ignore only 'table "XXX" does not exist' errors.
+            connection.rollback();
+        } finally {
+            stmt.close();
+        }
+    }
+
+    // Redshift does not support DROP TABLE IF EXISTS.
+    // Dropping part runs DROP TABLE and ignores errors.
+    @Override
+    public void replaceTable(String fromTable, JdbcSchema schema, String toTable) throws SQLException
+    {
+        Statement stmt = connection.createStatement();
+        try {
+            try {
+                StringBuilder sb = new StringBuilder();
+                sb.append("DROP TABLE ");
+                quoteIdentifierString(sb, toTable);
+                String sql = sb.toString();
+                executeUpdate(stmt, sql);
+            } catch (SQLException ex) {
+                // ignore errors.
+                // TODO here should ignore only 'table "XXX" does not exist' errors.
+                // rollback or comimt is required to recover failed transaction
+                connection.rollback();
+            }
+
+            {
+                StringBuilder sb = new StringBuilder();
+                sb.append("ALTER TABLE ");
+                quoteIdentifierString(sb, fromTable);
+                sb.append(" RENAME TO ");
+                quoteIdentifierString(sb, toTable);
+                String sql = sb.toString();
+                executeUpdate(stmt, sql);
+            }
+
+            connection.commit();
+        } catch (SQLException ex) {
+            connection.rollback();
+            throw ex;
+        } finally {
+            stmt.close();
+        }
+    }
+
+    @Override
+    protected String convertTypeName(String typeName)
+    {
+        // Redshift does not support TEXT type.
+        switch(typeName) {
+        case "CLOB":
+            return "VARCHAR(65535)";
+        case "TEXT":
+            return "VARCHAR(65535)";
+        case "BLOB":
+            return "BYTEA";
+        default:
+            return typeName;
+        }
+    }
+
+    public String buildCopySQLBeforeFrom(String tableName, JdbcSchema tableSchema)
+    {
+        StringBuilder sb = new StringBuilder();
+
+        sb.append("COPY ");
+        quoteIdentifierString(sb, tableName);
+        sb.append(" (");
+        for(int i=0; i < tableSchema.getCount(); i++) {
+            if(i != 0) { sb.append(", "); }
+            quoteIdentifierString(sb, tableSchema.getColumnName(i));
+        }
+        sb.append(")");
+
+        return sb.toString();
+    }
+
+    public void runCopy(String sql) throws SQLException
+    {
+        Statement stmt = connection.createStatement();
+        try {
+            stmt.executeUpdate(sql);
+        } finally {
+            stmt.close();
+        }
+    }
+}

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

@@ -0,0 +1,40 @@
+package org.embulk.output.redshift;
+
+import java.util.Properties;
+import java.sql.Driver;
+import java.sql.Connection;
+import java.sql.SQLException;
+import org.embulk.output.jdbc.JdbcOutputConnector;
+import org.embulk.output.jdbc.JdbcOutputConnection;
+
+public class RedshiftOutputConnector
+        implements JdbcOutputConnector
+{
+    private static final Driver driver = new org.postgresql.Driver();
+
+    private final String url;
+    private final Properties properties;
+    private final String schemaName;
+
+    public RedshiftOutputConnector(String url, Properties properties, String schemaName)
+    {
+        this.url = url;
+        this.properties = properties;
+        this.schemaName = schemaName;
+    }
+
+    @Override
+    public RedshiftOutputConnection connect(boolean autoCommit) throws SQLException
+    {
+        Connection c = driver.connect(url, properties);
+        try {
+            RedshiftOutputConnection con = new RedshiftOutputConnection(c, schemaName, autoCommit);
+            c = null;
+            return con;
+        } finally {
+            if (c != null) {
+                c.close();
+            }
+        }
+    }
+}

metadata

@@ -0,0 +1,67 @@
+--- !ruby/object:Gem::Specification
+name: embulk-output-redshift
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- FURUHASHI Sadayuki
+autorequire:
+bindir: bin
+cert_chain: []
+date: 2015-02-16 00:00:00.000000000 Z
+dependencies: []
+description: JDBC output plugin is an Embulk plugin that loads records to JDBC read by any input plugins. Search the input plugins by "embulk-input" keyword.
+email:
+- frsyuki@users.sourceforge.jp
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- README.md
+- build.gradle
+- lib/embulk/output/redshift.rb
+- src/main/java/org/embulk/output/RedshiftOutputPlugin.java
+- src/main/java/org/embulk/output/redshift/RedshiftCopyBatchInsert.java
+- src/main/java/org/embulk/output/redshift/RedshiftOutputConnection.java
+- src/main/java/org/embulk/output/redshift/RedshiftOutputConnector.java
+- classpath/aws-java-sdk-core-1.9.17.jar
+- classpath/aws-java-sdk-kms-1.9.17.jar
+- classpath/aws-java-sdk-s3-1.9.17.jar
+- classpath/aws-java-sdk-sts-1.9.17.jar
+- classpath/commons-codec-1.6.jar
+- classpath/commons-logging-1.1.3.jar
+- classpath/embulk-output-jdbc-0.1.0.jar
+- classpath/embulk-output-postgresql-0.1.0.jar
+- classpath/embulk-output-redshift-0.1.0.jar
+- classpath/httpclient-4.3.4.jar
+- classpath/httpcore-4.3.2.jar
+- classpath/jna-4.1.0.jar
+- classpath/jna-platform-4.1.0.jar
+- classpath/postgresql-9.4-1200-jdbc41.jar
+- classpath/slf4j-simple-1.7.7.jar
+- classpath/waffle-jna-1.7.jar
+homepage: https://github.com/embulk/embulk-output-jdbc
+licenses:
+- Apache 2.0
+metadata: {}
+post_install_message:
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - '>='
+    - !ruby/object:Gem::Version
+      version: '0'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - '>='
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubyforge_project:
+rubygems_version: 2.1.9
+signing_key:
+specification_version: 4
+summary: JDBC output plugin for Embulk
+test_files: []