embulk-output-postgresql 0.4.1 → 0.4.2

data/src/main/java/org/embulk/output/postgresql/PostgreSQLCopyBatchInsert.java

@@ -1,73 +1,71 @@
-package org.embulk.output.postgresql;
-
-import java.io.File;
-import java.io.IOException;
-import java.io.FileInputStream;
-import java.sql.Connection;
-import java.sql.SQLException;
-import org.slf4j.Logger;
-import org.postgresql.copy.CopyManager;
-import org.postgresql.core.BaseConnection;
-import org.embulk.spi.Exec;
-import org.embulk.output.jdbc.JdbcSchema;
-
-public class PostgreSQLCopyBatchInsert
-        extends AbstractPostgreSQLCopyBatchInsert
-{
-    private final Logger logger = Exec.getLogger(PostgreSQLCopyBatchInsert.class);
-    private final PostgreSQLOutputConnector connector;
-
-    private PostgreSQLOutputConnection connection = null;
-    private CopyManager copyManager = null;
-    private String copySql = null;
-    private long totalRows;
-
-    public PostgreSQLCopyBatchInsert(PostgreSQLOutputConnector connector) throws IOException, SQLException
-    {
-        super();
-        this.connector = connector;
-    }
-
-    @Override
-    public void prepare(String loadTable, JdbcSchema insertSchema) throws SQLException
-    {
-        this.connection = connector.connect(true);
-        this.copySql = connection.buildCopySql(loadTable, insertSchema);
-        this.copyManager = connection.newCopyManager();
-        logger.info("Copy SQL: "+copySql);
-    }
-
-    @Override
-    public void flush() throws IOException, SQLException
-    {
-        File file = closeCurrentFile();  // flush buffered data in writer
-
-        logger.info(String.format("Loading %,d rows (%,d bytes)", batchRows, file.length()));
-        long startTime = System.currentTimeMillis();
-        FileInputStream in = new FileInputStream(file);
-        try {
-            // TODO check age of connection and call isValid if it's old and reconnect if it's invalid
-            copyManager.copyIn(copySql, in);
-        } finally {
-            in.close();
-        }
-        double seconds = (System.currentTimeMillis() - startTime) / 1000.0;
-
-        totalRows += batchRows;
-        batchRows = 0;
-        logger.info(String.format("> %.2f seconds (loaded %,d rows in total)", seconds, totalRows));
-
-        openNewFile();
-        file.delete();
-    }
-
-    @Override
-    public void close() throws IOException, SQLException
-    {
-        closeCurrentFile().delete();
-        if (connection != null) {
-            connection.close();
-            connection = null;
-        }
-    }
-}
+package org.embulk.output.postgresql;
+
+import java.io.File;
+import java.io.IOException;
+import java.io.FileInputStream;
+import java.sql.SQLException;
+import org.slf4j.Logger;
+import org.postgresql.copy.CopyManager;
+import org.embulk.spi.Exec;
+import org.embulk.output.jdbc.JdbcSchema;
+
+public class PostgreSQLCopyBatchInsert
+        extends AbstractPostgreSQLCopyBatchInsert
+{
+    private final Logger logger = Exec.getLogger(PostgreSQLCopyBatchInsert.class);
+    private final PostgreSQLOutputConnector connector;
+
+    private PostgreSQLOutputConnection connection = null;
+    private CopyManager copyManager = null;
+    private String copySql = null;
+    private long totalRows;
+
+    public PostgreSQLCopyBatchInsert(PostgreSQLOutputConnector connector) throws IOException, SQLException
+    {
+        super();
+        this.connector = connector;
+    }
+
+    @Override
+    public void prepare(String loadTable, JdbcSchema insertSchema) throws SQLException
+    {
+        this.connection = connector.connect(true);
+        this.copySql = connection.buildCopySql(loadTable, insertSchema);
+        this.copyManager = connection.newCopyManager();
+        logger.info("Copy SQL: "+copySql);
+    }
+
+    @Override
+    public void flush() throws IOException, SQLException
+    {
+        File file = closeCurrentFile();  // flush buffered data in writer
+
+        logger.info(String.format("Loading %,d rows (%,d bytes)", batchRows, file.length()));
+        long startTime = System.currentTimeMillis();
+        FileInputStream in = new FileInputStream(file);
+        try {
+            // TODO check age of connection and call isValid if it's old and reconnect if it's invalid
+            copyManager.copyIn(copySql, in);
+        } finally {
+            in.close();
+        }
+        double seconds = (System.currentTimeMillis() - startTime) / 1000.0;
+
+        totalRows += batchRows;
+        batchRows = 0;
+        logger.info(String.format("> %.2f seconds (loaded %,d rows in total)", seconds, totalRows));
+
+        openNewFile();
+        file.delete();
+    }
+
+    @Override
+    public void close() throws IOException, SQLException
+    {
+        closeCurrentFile().delete();
+        if (connection != null) {
+            connection.close();
+            connection = null;
+        }
+    }
+}

data/src/main/java/org/embulk/output/postgresql/PostgreSQLOutputConnection.java

@@ -1,157 +1,156 @@
-package org.embulk.output.postgresql;
-
-import java.util.List;
-import java.sql.Connection;
-import java.sql.SQLException;
-import java.sql.Statement;
-import org.postgresql.copy.CopyManager;
-import org.postgresql.core.BaseConnection;
-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 PostgreSQLOutputConnection
-        extends JdbcOutputConnection
-{
-    public PostgreSQLOutputConnection(Connection connection, String schemaName, boolean autoCommit)
-            throws SQLException
-    {
-        super(connection, schemaName);
-        connection.setAutoCommit(autoCommit);
-    }
-
-    public String buildCopySql(String toTable, JdbcSchema toTableSchema)
-    {
-        StringBuilder sb = new StringBuilder();
-
-        sb.append("COPY ");
-        quoteIdentifierString(sb, toTable);
-        sb.append(" (");
-        for (int i=0; i < toTableSchema.getCount(); i++) {
-            if (i != 0) { sb.append(", "); }
-            quoteIdentifierString(sb, toTableSchema.getColumnName(i));
-        }
-        sb.append(") ");
-        sb.append("FROM STDIN");
-
-        return sb.toString();
-    }
-
-    public CopyManager newCopyManager() throws SQLException
-    {
-        return new CopyManager((BaseConnection) connection);
-    }
-
-    @Override
-    protected String buildCollectMergeSql(List<String> fromTables, JdbcSchema schema, String toTable, List<String> mergeKeys) throws SQLException
-    {
-        StringBuilder sb = new StringBuilder();
-
-        sb.append("WITH updated AS (");
-        sb.append("UPDATE ");
-        quoteIdentifierString(sb, toTable);
-        sb.append(" SET ");
-        for (int i=0; i < schema.getCount(); i++) {
-            if (i != 0) { sb.append(", "); }
-            quoteIdentifierString(sb, schema.getColumnName(i));
-            sb.append(" = S.");
-            quoteIdentifierString(sb, schema.getColumnName(i));
-        }
-        sb.append(" FROM (");
-        for (int i=0; i < fromTables.size(); i++) {
-            if (i != 0) { sb.append(" UNION ALL "); }
-            sb.append("SELECT ");
-            for(int j=0; j < schema.getCount(); j++) {
-                if (j != 0) { sb.append(", "); }
-                quoteIdentifierString(sb, schema.getColumnName(j));
-            }
-            sb.append(" FROM ");
-            quoteIdentifierString(sb, fromTables.get(i));
-        }
-        sb.append(") S");
-        sb.append(" WHERE ");
-        for (int i=0; i < mergeKeys.size(); i++) {
-            if (i != 0) { sb.append(" AND "); }
-            quoteIdentifierString(sb, toTable);
-            sb.append(".");
-            quoteIdentifierString(sb, mergeKeys.get(i));
-            sb.append(" = ");
-            sb.append("S.");
-            quoteIdentifierString(sb, mergeKeys.get(i));
-        }
-        sb.append(" RETURNING ");
-        for (int i=0; i < mergeKeys.size(); i++) {
-            if (i != 0) { sb.append(", "); }
-            sb.append("S.");
-            quoteIdentifierString(sb, mergeKeys.get(i));
-        }
-        sb.append(") ");
-
-        sb.append("INSERT INTO ");
-        quoteIdentifierString(sb, toTable);
-        sb.append(" (");
-        for (int i=0; i < schema.getCount(); i++) {
-            if (i != 0) { sb.append(", "); }
-            quoteIdentifierString(sb, schema.getColumnName(i));
-        }
-        sb.append(") ");
-        sb.append("SELECT DISTINCT ON (");
-        for (int i=0; i < mergeKeys.size(); i++) {
-            if (i != 0) { sb.append(", "); }
-            quoteIdentifierString(sb, mergeKeys.get(i));
-        }
-        sb.append(") * FROM (");
-        for (int i=0; i < fromTables.size(); i++) {
-            if (i != 0) { sb.append(" UNION ALL "); }
-            sb.append("SELECT ");
-            for(int j=0; j < schema.getCount(); j++) {
-                if (j != 0) { sb.append(", "); }
-                quoteIdentifierString(sb, schema.getColumnName(j));
-            }
-            sb.append(" FROM ");
-            quoteIdentifierString(sb, fromTables.get(i));
-        }
-        sb.append(") S ");
-        sb.append("WHERE NOT EXISTS (");
-        sb.append("SELECT 1 FROM updated WHERE ");
-        for (int i=0; i < mergeKeys.size(); i++) {
-            if (i != 0) { sb.append(" AND "); }
-            sb.append("S.");
-            quoteIdentifierString(sb, mergeKeys.get(i));
-            sb.append(" = ");
-            sb.append("updated.");
-            quoteIdentifierString(sb, mergeKeys.get(i));
-        }
-        sb.append(") ");
-
-        return sb.toString();
-    }
-
-    protected void collectReplaceView(List<String> fromTables, JdbcSchema schema, String toTable) throws SQLException
-    {
-        Statement stmt = connection.createStatement();
-        try {
-            String sql = buildCollectInsertSql(fromTables, schema, toTable);
-            executeUpdate(stmt, sql);
-            commitIfNecessary(connection);
-        } catch (SQLException ex) {
-            throw safeRollback(connection, ex);
-        } finally {
-            stmt.close();
-        }
-    }
-
-    @Override
-    protected String buildColumnTypeName(JdbcColumn c)
-    {
-        switch(c.getSimpleTypeName()) {
-        case "CLOB":
-            return "TEXT";
-        case "BLOB":
-            return "BYTEA";
-        default:
-            return super.buildColumnTypeName(c);
-        }
-    }
-}
+package org.embulk.output.postgresql;
+
+import java.util.List;
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.sql.Statement;
+import org.postgresql.copy.CopyManager;
+import org.postgresql.core.BaseConnection;
+import org.embulk.output.jdbc.JdbcOutputConnection;
+import org.embulk.output.jdbc.JdbcColumn;
+import org.embulk.output.jdbc.JdbcSchema;
+
+public class PostgreSQLOutputConnection
+        extends JdbcOutputConnection
+{
+    public PostgreSQLOutputConnection(Connection connection, String schemaName, boolean autoCommit)
+            throws SQLException
+    {
+        super(connection, schemaName);
+        connection.setAutoCommit(autoCommit);
+    }
+
+    public String buildCopySql(String toTable, JdbcSchema toTableSchema)
+    {
+        StringBuilder sb = new StringBuilder();
+
+        sb.append("COPY ");
+        quoteIdentifierString(sb, toTable);
+        sb.append(" (");
+        for (int i=0; i < toTableSchema.getCount(); i++) {
+            if (i != 0) { sb.append(", "); }
+            quoteIdentifierString(sb, toTableSchema.getColumnName(i));
+        }
+        sb.append(") ");
+        sb.append("FROM STDIN");
+
+        return sb.toString();
+    }
+
+    public CopyManager newCopyManager() throws SQLException
+    {
+        return new CopyManager((BaseConnection) connection);
+    }
+
+    @Override
+    protected String buildCollectMergeSql(List<String> fromTables, JdbcSchema schema, String toTable, List<String> mergeKeys) throws SQLException
+    {
+        StringBuilder sb = new StringBuilder();
+
+        sb.append("WITH updated AS (");
+        sb.append("UPDATE ");
+        quoteIdentifierString(sb, toTable);
+        sb.append(" SET ");
+        for (int i=0; i < schema.getCount(); i++) {
+            if (i != 0) { sb.append(", "); }
+            quoteIdentifierString(sb, schema.getColumnName(i));
+            sb.append(" = S.");
+            quoteIdentifierString(sb, schema.getColumnName(i));
+        }
+        sb.append(" FROM (");
+        for (int i=0; i < fromTables.size(); i++) {
+            if (i != 0) { sb.append(" UNION ALL "); }
+            sb.append("SELECT ");
+            for(int j=0; j < schema.getCount(); j++) {
+                if (j != 0) { sb.append(", "); }
+                quoteIdentifierString(sb, schema.getColumnName(j));
+            }
+            sb.append(" FROM ");
+            quoteIdentifierString(sb, fromTables.get(i));
+        }
+        sb.append(") S");
+        sb.append(" WHERE ");
+        for (int i=0; i < mergeKeys.size(); i++) {
+            if (i != 0) { sb.append(" AND "); }
+            quoteIdentifierString(sb, toTable);
+            sb.append(".");
+            quoteIdentifierString(sb, mergeKeys.get(i));
+            sb.append(" = ");
+            sb.append("S.");
+            quoteIdentifierString(sb, mergeKeys.get(i));
+        }
+        sb.append(" RETURNING ");
+        for (int i=0; i < mergeKeys.size(); i++) {
+            if (i != 0) { sb.append(", "); }
+            sb.append("S.");
+            quoteIdentifierString(sb, mergeKeys.get(i));
+        }
+        sb.append(") ");
+
+        sb.append("INSERT INTO ");
+        quoteIdentifierString(sb, toTable);
+        sb.append(" (");
+        for (int i=0; i < schema.getCount(); i++) {
+            if (i != 0) { sb.append(", "); }
+            quoteIdentifierString(sb, schema.getColumnName(i));
+        }
+        sb.append(") ");
+        sb.append("SELECT DISTINCT ON (");
+        for (int i=0; i < mergeKeys.size(); i++) {
+            if (i != 0) { sb.append(", "); }
+            quoteIdentifierString(sb, mergeKeys.get(i));
+        }
+        sb.append(") * FROM (");
+        for (int i=0; i < fromTables.size(); i++) {
+            if (i != 0) { sb.append(" UNION ALL "); }
+            sb.append("SELECT ");
+            for(int j=0; j < schema.getCount(); j++) {
+                if (j != 0) { sb.append(", "); }
+                quoteIdentifierString(sb, schema.getColumnName(j));
+            }
+            sb.append(" FROM ");
+            quoteIdentifierString(sb, fromTables.get(i));
+        }
+        sb.append(") S ");
+        sb.append("WHERE NOT EXISTS (");
+        sb.append("SELECT 1 FROM updated WHERE ");
+        for (int i=0; i < mergeKeys.size(); i++) {
+            if (i != 0) { sb.append(" AND "); }
+            sb.append("S.");
+            quoteIdentifierString(sb, mergeKeys.get(i));
+            sb.append(" = ");
+            sb.append("updated.");
+            quoteIdentifierString(sb, mergeKeys.get(i));
+        }
+        sb.append(") ");
+
+        return sb.toString();
+    }
+
+    protected void collectReplaceView(List<String> fromTables, JdbcSchema schema, String toTable) throws SQLException
+    {
+        Statement stmt = connection.createStatement();
+        try {
+            String sql = buildCollectInsertSql(fromTables, schema, toTable);
+            executeUpdate(stmt, sql);
+            commitIfNecessary(connection);
+        } catch (SQLException ex) {
+            throw safeRollback(connection, ex);
+        } finally {
+            stmt.close();
+        }
+    }
+
+    @Override
+    protected String buildColumnTypeName(JdbcColumn c)
+    {
+        switch(c.getSimpleTypeName()) {
+        case "CLOB":
+            return "TEXT";
+        case "BLOB":
+            return "BYTEA";
+        default:
+            return super.buildColumnTypeName(c);
+        }
+    }
+}