embulk-output-jdbc 0.1.0

data/src/main/java/org/embulk/output/jdbc/JdbcOutputConnector.java

@@ -0,0 +1,8 @@
+package org.embulk.output.jdbc;
+
+import java.sql.SQLException;
+
+public interface JdbcOutputConnector
+{
+    public JdbcOutputConnection connect(boolean autoCommit) throws SQLException;
+}

data/src/main/java/org/embulk/output/jdbc/JdbcSchema.java

@@ -0,0 +1,37 @@
+package org.embulk.output.jdbc;
+
+import java.util.List;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonValue;
+
+public class JdbcSchema
+{
+    private List<JdbcColumn> columns;
+
+    @JsonCreator
+    public JdbcSchema(List<JdbcColumn> columns)
+    {
+        this.columns = columns;
+    }
+
+    @JsonValue
+    public List<JdbcColumn> getColumns()
+    {
+        return columns;
+    }
+
+    public int getCount()
+    {
+        return columns.size();
+    }
+
+    public JdbcColumn getColumn(int i)
+    {
+        return columns.get(i);
+    }
+
+    public String getColumnName(int i)
+    {
+        return columns.get(i).getName();
+    }
+}

data/src/main/java/org/embulk/output/jdbc/JdbcUtils.java

@@ -0,0 +1,155 @@
+package org.embulk.output.jdbc;
+
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import org.slf4j.Logger;
+import org.embulk.spi.Exec;
+
+public class JdbcUtils
+{
+    public final Logger logger = Exec.getLogger(JdbcUtils.class.getName());
+
+    private static String[] SEARCH_STRING_SPECIAL_CHARS = new String[] { "_", "%" };
+
+    public static String escapeSearchString(String searchString, String escapeString)
+    {
+        if (searchString != null && escapeString != null) {
+            // First of all, escape escapeString '\' itself: '\' -> '\\'
+            searchString = searchString.replaceAll(Pattern.quote(escapeString),
+                    Matcher.quoteReplacement(escapeString + escapeString));
+            for (String specialChar : SEARCH_STRING_SPECIAL_CHARS) {
+                if (specialChar.equals(escapeString)) {
+                    throw new IllegalArgumentException("Special char " + specialChar + " cannot be an escape char");
+                }
+                searchString = searchString.replaceAll(Pattern.quote(specialChar),
+                        Matcher.quoteReplacement(escapeString + specialChar));
+            }
+        }
+        return searchString;
+    }
+
+    private Class<?> connectionClass;
+
+    // Connection.isValid() is available from Java 1.6 + JDBC4.
+    //private boolean hasIsValid;
+
+    // Connection.setNetworkTimeout() is available from Java 1.7 + JDBC4.
+    //private boolean hasSetNetworkTimeout;
+    //private Method setNetworkTimeoutMethod;
+
+    public JdbcUtils(Class<?> connectionClass)
+    {
+        this.connectionClass = connectionClass;
+
+        //configureSetNetworkTimeout();
+    }
+
+    public int executeUpdateWithSqlLogging(Statement stmt, String sql) throws SQLException
+    {
+        logger.info("SQL: " + sql);
+        long startTime = System.currentTimeMillis();
+        int count = stmt.executeUpdate(sql);
+        double seconds = (System.currentTimeMillis() - startTime) / 1000.0;
+        if (count == 0) {
+            logger.info(String.format("> %.2f seconds", seconds));
+        } else {
+            logger.info(String.format("> %.2f seconds (%,d rows)", seconds, count));
+        }
+        return count;
+    }
+
+    //private void configureSetNetworkTimeout() {
+    //    try {
+    //        Method m = connectionClass.getMethod("setNetworkTimeout", Executor.class, int.class);
+    //        if (isCallableMethod(m)) {
+    //            setNetworkTimeoutMethod = m;
+    //            hasSetNetworkTimeout = true;
+    //        }
+    //    } catch (SecurityException ex) {
+    //    } catch (NoSuchMethodException ex) {
+    //    }
+    //}
+
+    //private boolean isCallableMethod(Method m) {
+    //    int modifiers = m.getModifiers();
+    //    if (Modifier.isAbstract(modifiers)) {
+    //        // Method.invoke throws java.lang.AbstractMethodError if it's an
+    //        // abstract method. Applications can't catch AbstractMethodError.
+    //        return false;
+    //    }
+    //    if (!Modifier.isPublic(modifiers)) {
+    //        // we can only call public methods
+    //        return false;
+    //    }
+    //    return true;
+    //}
+
+    // PostgreSQL JDBC driver implements isValid() method. But the
+    // implementation throws following exception:
+    // "java.io.IOException: Method org.postgresql.jdbc4.Jdbc4Connection.isValid(int) is not yet implemented."
+    //
+    // So, checking mechanism doesn't work at all.
+    // Thus here just runs "SELECT 1" to check connectivity.
+    //
+    //public boolean isValidConnection(Connection connection, int timeout) throws SQLException
+    //{
+    //    Statement stmt = connection.createStatement();
+    //    try {
+    //        stmt.executeQuery("SELECT 1").close();
+    //        return true;
+    //    } catch (SQLException ex) {
+    //        return false;
+    //    } finally {
+    //        stmt.close();
+    //    }
+    //}
+
+    //public void setNetworkTimeout(Connection connection,
+    //        Executor executor, int milliseconds) throws SQLException {
+    //    Throwable exception = null;
+    //    if (hasSetNetworkTimeout) {
+    //        try {
+    //            setNetworkTimeoutMethod.invoke(connection, executor, milliseconds);
+    //            return;
+    //
+    //        } catch (IllegalArgumentException ex) {
+    //            // ignore error
+    //            LOG.warn("Connection.setNetworkTimeout failed due to IllegalArgumentException.");
+    //            exception = ex;
+    //
+    //        } catch (IllegalAccessException ex) {
+    //            // ignore error
+    //            LOG.warn("Connection.setNetworkTimeout failed due to IllegalAccessException.");
+    //            exception = ex;
+    //
+    //        } catch (InvocationTargetException ex) {
+    //            //Throwable cause = ex.getTargetException();
+    //            //if (cause instanceof SQLException) {
+    //            //    throw (SQLException) cause;
+    //            //} else if (cause instanceof RuntimeException) {
+    //            //    throw (RuntimeException) cause;
+    //            //} else if (cause instanceof Error) {
+    //            //    throw (Error) cause;
+    //            //} else {
+    //            //    throw new SQLException(cause);
+    //            //}
+    //            exception = ex.getTargetException();
+    //            // It's safe to ignore exceptions.
+    //        }
+    //
+    //        hasSetNetworkTimeout = false;
+    //    }
+    //
+    //    if (exception != null) {
+    //        LOG.warn("Connection.setNetworkTimeout is not available: "+exception);
+    //    } else {
+    //        LOG.warn("Connection.setNetworkTimeout is not available.");
+    //    }
+    //    // TODO any substitute implementations?
+    //}
+}
+

data/src/main/java/org/embulk/output/jdbc/RetryExecutor.java

@@ -0,0 +1,105 @@
+package org.embulk.output.jdbc;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+
+public class RetryExecutor
+{
+    public static RetryExecutor retryExecutor()
+    {
+        return new RetryExecutor();
+    }
+
+    public static abstract class IdempotentOperation<T> implements Callable<T>
+    {
+        public abstract T call() throws Exception;
+
+        public void onRetry(Throwable exception, int retryCount, int retryLimit, int retryWait)
+        { }
+
+        public void onGiveup(Throwable firstException, Throwable lastException)
+        { }
+
+        public abstract boolean isRetryableException(Throwable exception);
+    }
+
+    private int retryLimit = 3;
+    private int initialRetryWait = 500;
+    private int maxRetryWait = 30*60*1000;
+
+    private RetryExecutor()
+    { }
+
+    public RetryExecutor setRetryLimit(int count)
+    {
+        this.retryLimit = count;
+        return this;
+    }
+
+    public RetryExecutor setInitialRetryWait(int msec)
+    {
+        this.initialRetryWait = msec;
+        return this;
+    }
+
+    public RetryExecutor setMaxRetryWait(int msec)
+    {
+        this.maxRetryWait = msec;
+        return this;
+    }
+
+    public <T> T runInterruptible(IdempotentOperation<T> op) throws InterruptedException, ExecutionException
+    {
+        return run(op, true);
+    }
+
+    public <T> T run(IdempotentOperation<T> op) throws ExecutionException
+    {
+        try {
+            return run(op, false);
+        } catch (InterruptedException ex) {
+            throw new ExecutionException("Unexpected interruption", ex);
+        }
+    }
+
+    private <T> T run(IdempotentOperation<T> op, boolean interruptible)
+            throws InterruptedException, ExecutionException
+    {
+        int retryWait = initialRetryWait;
+        int retryCount = 0;
+
+        Throwable firstException = null;
+
+        while(true) {
+            try {
+                return op.call();
+            } catch (Throwable exception) {
+                if (firstException == null) {
+                    firstException = exception;
+                }
+                if (!op.isRetryableException(exception) || retryCount >= retryLimit) {
+                    op.onGiveup(firstException, exception);
+                    throw new ExecutionException(firstException);
+                }
+
+                retryCount++;
+                op.onRetry(exception, retryCount, retryLimit, retryWait);
+
+                try {
+                    Thread.sleep(retryWait);
+                } catch (InterruptedException ex) {
+                    if (interruptible) {
+                        throw ex;
+                    }
+                }
+
+                retryWait *= 2;  // exponential back-off
+
+                if (retryWait > maxRetryWait) {
+                    retryWait = maxRetryWait;
+                }
+            }
+        }
+    }
+}
+

data/src/main/java/org/embulk/output/jdbc/StandardBatchInsert.java

@@ -0,0 +1,180 @@
+package org.embulk.output.jdbc;
+
+import java.io.IOException;
+import java.math.BigDecimal;
+import java.sql.PreparedStatement;
+import java.sql.SQLException;
+import java.sql.Date;
+import java.sql.Time;
+import java.sql.Timestamp;
+import org.slf4j.Logger;
+import org.embulk.spi.Exec;
+
+public class StandardBatchInsert
+        implements BatchInsert
+{
+    private final Logger logger = Exec.getLogger(StandardBatchInsert.class);
+
+    private final JdbcOutputConnector connector;
+
+    private JdbcOutputConnection connection;
+    private PreparedStatement batch;
+    private int index;
+    private int batchWeight;
+    private int batchRows;
+    private long totalRows;
+
+    public StandardBatchInsert(JdbcOutputConnector connector) throws IOException, SQLException
+    {
+        this.connector = connector;
+    }
+
+    public void prepare(String loadTable, JdbcSchema insertSchema) throws SQLException
+    {
+        this.connection = connector.connect(true);
+        this.batch = connection.prepareInsertSql(loadTable, insertSchema);
+        this.index = 1;  // PreparedStatement index begings from 1
+        this.batchRows = 0;
+        this.totalRows = 0;
+        batch.clearBatch();
+    }
+
+    public int getBatchWeight()
+    {
+        return batchWeight;
+    }
+
+    public void add() throws IOException, SQLException
+    {
+        batch.addBatch();
+        index = 1;  // PreparedStatement index begins from 1
+        batchRows++;
+        batchWeight += 32;  // add weight as overhead of each rows
+    }
+
+    public void close() throws IOException, SQLException
+    {
+        // caller should close the connection
+    }
+
+    public void flush() throws IOException, SQLException
+    {
+        logger.info(String.format("Loading %,d rows", batchRows));
+        long startTime = System.currentTimeMillis();
+        batch.executeBatch();  // here can't use returned value because MySQL Connector/J returns SUCCESS_NO_INFO as a batch result
+        double seconds = (System.currentTimeMillis() - startTime) / 1000.0;
+
+        totalRows += batchRows;
+        logger.info(String.format("> %.2f seconds (loaded %,d rows in total)", seconds, totalRows));
+        batch.clearBatch();
+        batchRows = 0;
+        batchWeight = 0;
+    }
+
+    public void finish() throws IOException, SQLException
+    {
+        if (getBatchWeight() != 0) {
+            flush();
+        }
+    }
+
+    public void setNull(int sqlType) throws IOException, SQLException
+    {
+        batch.setNull(index, sqlType);
+        nextColumn(0);
+    }
+
+    public void setBoolean(boolean v) throws IOException, SQLException
+    {
+        batch.setBoolean(index, v);
+        nextColumn(1);
+    }
+
+    public void setByte(byte v) throws IOException, SQLException
+    {
+        batch.setByte(index, v);
+        nextColumn(1);
+    }
+
+    public void setShort(short v) throws IOException, SQLException
+    {
+        batch.setShort(index, v);
+        nextColumn(2);
+    }
+
+    public void setInt(int v) throws IOException, SQLException
+    {
+        batch.setInt(index, v);
+        nextColumn(4);
+    }
+
+    public void setLong(long v) throws IOException, SQLException
+    {
+        batch.setLong(index, v);
+        nextColumn(8);
+    }
+
+    public void setFloat(float v) throws IOException, SQLException
+    {
+        batch.setFloat(index, v);
+        nextColumn(4);
+    }
+
+    public void setDouble(double v) throws IOException, SQLException
+    {
+        batch.setDouble(index, v);
+        nextColumn(8);
+    }
+
+    public void setBigDecimal(BigDecimal v) throws IOException, SQLException
+    {
+        // use estimated number of necessary bytes + 8 byte for the weight
+        // assuming one place needs 4 bits. ceil(v.precision() / 2.0) + 8
+        batch.setBigDecimal(index, v);
+        nextColumn((v.precision() & ~2) / 2 + 8);
+    }
+
+    public void setString(String v) throws IOException, SQLException
+    {
+        batch.setString(index, v);
+        // estimate all chracters use 2 bytes; almost enough for the worst case
+        nextColumn(v.length() * 2 + 4);
+    }
+
+    public void setNString(String v) throws IOException, SQLException
+    {
+        batch.setNString(index, v);
+        // estimate all chracters use 2 bytes; almost enough for the worst case
+        nextColumn(v.length() * 2 + 4);
+    }
+
+    public void setBytes(byte[] v) throws IOException, SQLException
+    {
+        batch.setBytes(index, v);
+        nextColumn(v.length + 4);
+    }
+
+    public void setSqlDate(Date v, int sqlType) throws IOException, SQLException
+    {
+        batch.setObject(index, v, sqlType);
+        nextColumn(32);
+    }
+
+    public void setSqlTime(Time v, int sqlType) throws IOException, SQLException
+    {
+        batch.setObject(index, v, sqlType);
+        nextColumn(32);
+    }
+
+    public void setSqlTimestamp(Timestamp v, int sqlType) throws IOException, SQLException
+    {
+        batch.setObject(index, v, sqlType);
+        nextColumn(32);
+    }
+
+    private void nextColumn(int weight)
+    {
+        index++;
+        batchWeight += weight + 4;  // add weight as overhead of each columns
+    }
+}