embulk-output-snowflake 0.1.1 → 0.2.0

data/src/main/java/org/embulk/output/snowflake/SnowflakeCopyBatchInsert.java

@@ -0,0 +1,373 @@
+package org.embulk.output.snowflake;
+
+import org.embulk.output.jdbc.BatchInsert;
+import org.embulk.output.jdbc.JdbcOutputConnector;
+import org.embulk.output.jdbc.JdbcSchema;
+import org.embulk.output.jdbc.TableIdentifier;
+import org.embulk.spi.time.Timestamp;
+import org.slf4j.Logger;
+
+import java.io.*;
+import java.math.BigDecimal;
+import java.nio.charset.Charset;
+import java.sql.SQLException;
+import java.util.*;
+import java.util.concurrent.*;
+import java.util.zip.GZIPOutputStream;
+import org.slf4j.LoggerFactory;
+
+public class SnowflakeCopyBatchInsert implements BatchInsert {
+    private final Logger logger = LoggerFactory.getLogger(SnowflakeCopyBatchInsert.class);
+    private final JdbcOutputConnector connector;
+    protected static final Charset FILE_CHARSET = Charset.forName("UTF-8");
+    private final ExecutorService executorService;
+    private final StageIdentifier stageIdentifier;
+    private final boolean deleteStageFile;
+
+    protected static final String nullString = "\\N";
+    protected static final String newLineString = "\n";
+    protected static final String delimiterString = "\t";
+
+    private SnowflakeOutputConnection connection = null;
+    private TableIdentifier tableIdentifier = null;
+    protected File currentFile;
+    protected BufferedWriter writer;
+    protected int index;
+    protected int batchRows;
+    private long totalRows;
+    private int fileCount;
+    private List<Future<Void>> uploadAndCopyFutures;
+
+    public SnowflakeCopyBatchInsert(JdbcOutputConnector connector, StageIdentifier stageIdentifier,
+                                    boolean deleteStageFile) throws IOException {
+        this.index = 0;
+        openNewFile();
+        this.connector = connector;
+        this.stageIdentifier = stageIdentifier;
+        this.executorService = Executors.newCachedThreadPool();
+        this.deleteStageFile = deleteStageFile;
+        this.uploadAndCopyFutures = new ArrayList();
+    }
+
+    @Override
+    public void prepare(TableIdentifier loadTable, JdbcSchema insertSchema) throws SQLException {
+        this.connection = (SnowflakeOutputConnection) connector.connect(true);
+        this.connection.runCreateStage(stageIdentifier);
+        this.tableIdentifier = loadTable;
+    }
+
+    private File createTempFile() throws IOException {
+        return File.createTempFile("embulk-output-snowflake-copy-", ".tsv.tmp");  // TODO configurable temporary file path
+    }
+
+    protected File openNewFile() throws IOException {
+        File newFile = createTempFile();
+        File oldFile = closeCurrentFile();
+        this.writer = openWriter(newFile);
+        currentFile = newFile;
+        return oldFile;
+    }
+
+    protected File closeCurrentFile() throws IOException {
+        if (writer != null) {
+            writer.close();
+            writer = null;
+        }
+        return currentFile;
+    }
+
+    protected BufferedWriter openWriter(File newFile) throws IOException {
+        // Snowflake supports gzip
+        return new BufferedWriter(
+                new OutputStreamWriter(
+                        new GZIPOutputStream(new FileOutputStream(newFile)),
+                        FILE_CHARSET)
+        );
+    }
+
+    public int getBatchWeight() {
+        long fsize = currentFile.length();
+        if (fsize > Integer.MAX_VALUE) {
+            return Integer.MAX_VALUE;
+        } else {
+            return (int) fsize;
+        }
+    }
+
+    public void add() throws IOException {
+        writer.write(newLineString);
+        batchRows++;
+        index = 0;
+    }
+
+    private void appendDelimiter() throws IOException {
+        if (index != 0) {
+            writer.write(delimiterString);
+        }
+        index++;
+    }
+
+    public void setNull(int sqlType) throws IOException {
+        appendDelimiter();
+        writer.write(nullString);
+    }
+
+    public void setBoolean(boolean v) throws IOException {
+        appendDelimiter();
+        writer.write(String.valueOf(v));
+    }
+
+    public void setByte(byte v) throws IOException {
+        appendDelimiter();
+        setEscapedString(String.valueOf(v));
+    }
+
+    public void setShort(short v) throws IOException {
+        appendDelimiter();
+        writer.write(String.valueOf(v));
+    }
+
+    public void setInt(int v) throws IOException {
+        appendDelimiter();
+        writer.write(String.valueOf(v));
+    }
+
+    public void setLong(long v) throws IOException {
+        appendDelimiter();
+        writer.write(String.valueOf(v));
+    }
+
+    public void setFloat(float v) throws IOException {
+        appendDelimiter();
+        writer.write(String.valueOf(v));
+    }
+
+    public void setDouble(double v) throws IOException {
+        appendDelimiter();
+        writer.write(String.valueOf(v));
+    }
+
+    public void setBigDecimal(BigDecimal v) throws IOException {
+        appendDelimiter();
+        writer.write(String.valueOf(v));
+    }
+
+    public void setString(String v) throws IOException {
+        appendDelimiter();
+        setEscapedString(v);
+    }
+
+    public void setNString(String v) throws IOException {
+        appendDelimiter();
+        setEscapedString(v);
+    }
+
+    public void setBytes(byte[] v) throws IOException {
+        appendDelimiter();
+        setEscapedString(String.valueOf(v));
+    }
+
+    public void setSqlDate(Timestamp v, Calendar cal) throws IOException {
+        appendDelimiter();
+        cal.setTimeInMillis(v.getEpochSecond() * 1000);
+        String f = String.format(Locale.ENGLISH, "%04d-%02d-%02d",
+                cal.get(Calendar.YEAR),
+                cal.get(Calendar.MONTH) + 1,
+                cal.get(Calendar.DAY_OF_MONTH));
+        writer.write(f);
+    }
+
+    public void setSqlTime(Timestamp v, Calendar cal) throws IOException {
+        appendDelimiter();
+        cal.setTimeInMillis(v.getEpochSecond() * 1000);
+        String f = String.format(Locale.ENGLISH, "%02d:%02d:%02d.%06d",
+                cal.get(Calendar.HOUR_OF_DAY),
+                cal.get(Calendar.MINUTE),
+                cal.get(Calendar.SECOND),
+                v.getNano() / 1000);
+        writer.write(f);
+    }
+
+    public void setSqlTimestamp(Timestamp v, Calendar cal) throws IOException {
+        appendDelimiter();
+        cal.setTimeInMillis(v.getEpochSecond() * 1000);
+        int zoneOffset = cal.get(Calendar.ZONE_OFFSET) / 1000 / 60;  // zone offset considering DST in minute
+        String offset;
+        if (zoneOffset >= 0) {
+            offset = String.format(Locale.ENGLISH, "+%02d%02d", zoneOffset / 60, zoneOffset % 60);
+        } else {
+            offset = String.format(Locale.ENGLISH, "-%02d%02d", -zoneOffset / 60, -zoneOffset % 60);
+        }
+        String f = String.format(Locale.ENGLISH, "%d-%02d-%02d %02d:%02d:%02d.%06d%s",
+                cal.get(Calendar.YEAR),
+                cal.get(Calendar.MONTH) + 1,
+                cal.get(Calendar.DAY_OF_MONTH),
+                cal.get(Calendar.HOUR_OF_DAY),
+                cal.get(Calendar.MINUTE),
+                cal.get(Calendar.SECOND),
+                v.getNano() / 1000,
+                offset);
+        writer.write(f);
+    }
+
+    private void setEscapedString(String v) throws IOException {
+        for (char c : v.toCharArray()) {
+            writer.write(escape(c));
+        }
+    }
+
+    @Override
+    public void flush() throws IOException, SQLException {
+        File file = closeCurrentFile();  // flush buffered data in writer
+
+        String snowflakeStageFileName = "embulk_snowflake_" + SnowflakeUtils.randomString(8);
+
+        UploadTask uploadTask = new UploadTask(file, batchRows, stageIdentifier, snowflakeStageFileName);
+        Future<Void> uploadFuture = executorService.submit(uploadTask);
+        uploadAndCopyFutures.add(uploadFuture);
+
+        CopyTask copyTask = new CopyTask(uploadFuture, snowflakeStageFileName);
+        uploadAndCopyFutures.add(executorService.submit(copyTask));
+
+        fileCount++;
+        totalRows += batchRows;
+        batchRows = 0;
+
+        openNewFile();
+    }
+
+    public void close() throws IOException, SQLException {
+        executorService.shutdownNow();
+
+        try {
+            executorService.awaitTermination(60, TimeUnit.SECONDS);
+        } catch (InterruptedException e) {
+        }
+
+        closeCurrentFile().delete();
+        if (connection != null) {
+            connection.close();
+            connection = null;
+        }
+    }
+
+    @Override
+    public void finish() throws IOException, SQLException
+    {
+        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 int[] getLastUpdateCounts()
+    {
+        // need not be implemented because SnowflakeCopyBatchInsert won't retry.
+        return new int[]{};
+    }
+
+    // Escape \, \n, \t, \r
+    // Remove \0
+    protected String escape(char c) {
+        switch (c) {
+            case '\\':
+                return "\\\\";
+            case '\n':
+                return "\\n";
+            case '\t':
+                return "\\t";
+            case '\r':
+                return "\\r";
+            case 0:
+                return "";
+            default:
+                return String.valueOf(c);
+        }
+    }
+
+    private class UploadTask implements Callable<Void> {
+        private final File file;
+        private final int batchRows;
+        private final String snowflakeStageFileName;
+        private final StageIdentifier stageIdentifier;
+
+        public UploadTask(File file, int batchRows, StageIdentifier stageIdentifier, String snowflakeStageFileName) {
+            this.file = file;
+            this.batchRows = batchRows;
+            this.snowflakeStageFileName = snowflakeStageFileName;
+            this.stageIdentifier = stageIdentifier;
+        }
+
+        public Void call() throws IOException, SQLException {
+            logger.info(String.format("Uploading file id %s to Snowflake (%,d bytes %,d rows)",
+                    snowflakeStageFileName, file.length(), batchRows));
+
+            try {
+                long startTime = System.currentTimeMillis();
+                // put file to snowflake internal storage
+                SnowflakeOutputConnection con = (SnowflakeOutputConnection) connector.connect(true);
+
+                FileInputStream fileInputStream = new FileInputStream(file);
+                con.runUploadFile(stageIdentifier, snowflakeStageFileName, fileInputStream);
+
+                double seconds = (System.currentTimeMillis() - startTime) / 1000.0;
+
+                logger.info(String.format("Uploaded file %s (%.2f seconds)", snowflakeStageFileName, seconds));
+            } finally {
+                file.delete();
+            }
+
+            return null;
+        }
+    }
+
+
+    private class CopyTask implements Callable<Void> {
+        private final Future<Void> uploadFuture;
+        private final String snowflakeStageFileName;
+
+        public CopyTask(Future<Void> uploadFuture, String snowflakeStageFileName) {
+            this.uploadFuture = uploadFuture;
+            this.snowflakeStageFileName = snowflakeStageFileName;
+        }
+
+        public Void call() throws SQLException, InterruptedException, ExecutionException {
+            try {
+                uploadFuture.get();
+
+                SnowflakeOutputConnection con = (SnowflakeOutputConnection) connector.connect(true);
+                try {
+                    logger.info("Running COPY from file {}", snowflakeStageFileName);
+
+                    long startTime = System.currentTimeMillis();
+                    con.runCopy(tableIdentifier, stageIdentifier, snowflakeStageFileName, delimiterString);
+
+                    double seconds = (System.currentTimeMillis() - startTime) / 1000.0;
+
+                    logger.info(String.format("Loaded file %s (%.2f seconds for COPY)", snowflakeStageFileName, seconds));
+
+                } finally {
+                    con.close();
+                }
+            } finally {
+                if (deleteStageFile) {
+                    connection.runDeleteStageFile(stageIdentifier, snowflakeStageFileName);
+                }
+            }
+
+            return null;
+        }
+    }
+}

data/src/main/java/org/embulk/output/snowflake/SnowflakeOutputConnection.java

@@ -1,51 +1,122 @@
 package org.embulk.output.snowflake;
 
+import net.snowflake.client.jdbc.SnowflakeConnection;
+import org.embulk.output.jdbc.JdbcColumn;
+import org.embulk.output.jdbc.JdbcOutputConnection;
+import org.embulk.output.jdbc.TableIdentifier;
+
+import java.io.FileInputStream;
 import java.sql.Connection;
-import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
 
-import org.embulk.output.jdbc.JdbcOutputConnection;
-import org.embulk.output.jdbc.JdbcUtils;
-import org.embulk.output.jdbc.TableIdentifier;
-
-public class SnowflakeOutputConnection
-        extends JdbcOutputConnection
-{
-    public SnowflakeOutputConnection(Connection connection, String schemaName)
-        throws SQLException
+public class SnowflakeOutputConnection extends JdbcOutputConnection {
+    public SnowflakeOutputConnection(Connection connection)
+            throws SQLException
     {
-        super(connection, schemaName);
+        super(connection, null);
     }
 
-    @Override
-    public boolean tableExists(TableIdentifier table) throws SQLException
+    public void runCopy(TableIdentifier tableIdentifier, StageIdentifier stageIdentifier, String filename, String delimiterString) throws SQLException
     {
-        String schemaName = JdbcUtils.escapeSearchString(table.getSchemaName(), connection.getMetaData().getSearchStringEscape());
-        String database = connection.getCatalog();
-        try (ResultSet rs = connection.getMetaData().getTables(database, schemaName, table.getTableName(), null)) {
-            return rs.next();
-        }
+        String sql = buildCopySQL(tableIdentifier, stageIdentifier, filename, delimiterString);
+        runUpdate(sql);
+    }
+    public void runCreateStage(StageIdentifier stageIdentifier) throws SQLException {
+        String sql = buildCreateStageSQL(stageIdentifier);
+        runUpdate(sql);
     }
 
-    @Override
-    public boolean tableExists(String tableName) throws SQLException
-    {
-        return tableExists(new TableIdentifier(connection.getCatalog(), schemaName, tableName));
+    public void runDropStage(StageIdentifier stageIdentifier) throws SQLException {
+        String sql = buildDropStageSQL(stageIdentifier);
+        runUpdate(sql);
     }
 
-    @Override
-    protected void setSearchPath(String schema) throws SQLException
+    public void runUploadFile(StageIdentifier stageIdentifier, String filename ,FileInputStream fileInputStream) throws SQLException{
+        connection.unwrap(SnowflakeConnection.class).uploadStream(stageIdentifier.getStageName(), stageIdentifier.getDestPrefix().or("/"),
+                fileInputStream, filename + ".csv.gz", false);
+    }
+
+    public void runDeleteStageFile(StageIdentifier stageIdentifier, String filename) throws SQLException{
+        String sql = buildDeleteStageFileSQL(stageIdentifier, filename);
+        runUpdate(sql);
+    }
+
+    protected void runUpdate(String sql) throws SQLException
     {
         Statement stmt = connection.createStatement();
         try {
-            String sql = "USE SCHEMA " + quoteIdentifierString(schema);
-            executeUpdate(stmt, sql);
-            commitIfNecessary(connection);
+            stmt.executeUpdate(sql);
         } finally {
             stmt.close();
         }
     }
 
+    @Override
+    protected String buildColumnTypeName(JdbcColumn c) {
+        switch(c.getSimpleTypeName()) {
+            case "CLOB":
+                return "VARCHAR(65535)";
+            default:
+                return super.buildColumnTypeName(c);
+        }
+    }
+
+    protected String buildCreateStageSQL(StageIdentifier stageIdentifier){
+        StringBuilder sb = new StringBuilder();
+        sb.append("CREATE STAGE IF NOT EXISTS ");
+        quoteStageIdentifier(sb, stageIdentifier);
+        sb.append(";");
+        return sb.toString();
+    }
+
+    protected String buildDropStageSQL(StageIdentifier stageIdentifier){
+        StringBuilder sb = new StringBuilder();
+        sb.append("DROP STAGE ");
+        quoteStageIdentifier(sb, stageIdentifier);
+        sb.append(";");
+        return sb.toString();
+    }
 
+    protected void quoteStageIdentifier(StringBuilder sb, StageIdentifier stageIdentifier){
+        sb.append(stageIdentifier.getDatabase());
+        sb.append(".");
+        sb.append(stageIdentifier.getSchemaName());
+        sb.append(".");
+        sb.append(stageIdentifier.getStageName());
+    }
+
+    protected String buildCopySQL(TableIdentifier tableIdentifier, StageIdentifier stageIdentifier, String snowflakeStageFileName, String delimiterString){
+        StringBuilder sb = new StringBuilder();
+        sb.append("COPY INTO ");
+        quoteTableIdentifier(sb, tableIdentifier);
+        sb.append(" FROM ");
+        quoteInternalStoragePath(sb, stageIdentifier, snowflakeStageFileName);
+        sb.append(" FILE_FORMAT = ( TYPE = CSV FIELD_DELIMITER = '");
+        sb.append(delimiterString);
+        sb.append("');");
+        return sb.toString();
+    }
+
+    protected String buildDeleteStageFileSQL(StageIdentifier stageIdentifier, String snowflakeStageFileName){
+        StringBuilder sb = new StringBuilder();
+        sb.append("REMOVE ");
+        quoteInternalStoragePath(sb, stageIdentifier, snowflakeStageFileName);
+        sb.append(';');
+        return sb.toString();
+    }
+
+    protected String quoteInternalStoragePath(StringBuilder sb, StageIdentifier stageIdentifier,
+                                              String snowflakeStageFileName){
+        sb.append("@");
+        quoteStageIdentifier(sb, stageIdentifier);
+        if (stageIdentifier.getDestPrefix().isPresent()){
+            sb.append("/");
+            sb.append(stageIdentifier.getDestPrefix().get());
+        }
+        sb.append("/");
+        sb.append(snowflakeStageFileName);
+        sb.append(".csv.gz");
+        return sb.toString();
+    }
 }