embulk-output-postgresql 0.2.4 → 0.3.0

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

@@ -1,73 +1,73 @@
-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.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;
+        }
+    }
+}

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

@@ -1,5 +1,6 @@
 package org.embulk.output.postgresql;
 
+import java.util.List;
 import java.sql.Connection;
 import java.sql.SQLException;
 import java.sql.Statement;
@@ -28,7 +29,7 @@ public class PostgreSQLOutputConnection
         quoteIdentifierString(sb, toTable);
         sb.append(" (");
         for (int i=0; i < toTableSchema.getCount(); i++) {
-            if(i != 0) { sb.append(", "); }
+            if (i != 0) { sb.append(", "); }
             quoteIdentifierString(sb, toTableSchema.getColumnName(i));
         }
         sb.append(") ");
@@ -43,66 +44,114 @@ public class PostgreSQLOutputConnection
     }
 
     @Override
-    protected String convertTypeName(String typeName)
-    {
-        switch(typeName) {
-        case "CLOB":
-            return "TEXT";
-        case "BLOB":
-            return "BYTEA";
-        default:
-            return typeName;
-        }
-    }
-
-    @Override
-    protected String buildPrepareUpsertSql(String toTable, JdbcSchema toTableSchema) throws SQLException
+    protected String buildCollectMergeSql(List<String> fromTables, JdbcSchema schema, String toTable, List<String> mergeKeys) throws SQLException
     {
         StringBuilder sb = new StringBuilder();
-        int size = toTableSchema.getCount();
-        String table = quoteIdentifierString(toTable);
-        int idx = 0;
 
-        sb.append("WITH upsert AS (UPDATE ").append(table).append(" SET ");
-
-        for (int i=0; i < size; i++) {
-            JdbcColumn c = toTableSchema.getColumn(i);
-            if (!c.isPrimaryKey()) {
-                if(idx != 0) { sb.append(", "); }
-                idx++;
-                quoteIdentifierString(sb, toTableSchema.getColumnName(i));
-                sb.append("=?");
+        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 ");
-        idx = 0;
-        for(int i=0; i < size; i++) {
-            JdbcColumn c = toTableSchema.getColumn(i);
-            if (c.isPrimaryKey()) {
-                if(idx != 0) { sb.append(" AND "); }
-                idx++;
-                quoteIdentifierString(sb, toTableSchema.getColumnName(i));
-                sb.append("=?");
-            }
+        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 true as result)");
-
-        sb.append(" INSERT INTO ").append(table).append(" (");
-        for (int i=0; i < size; i++) {
-            if(i != 0) { sb.append(", "); }
-            quoteIdentifierString(sb, toTableSchema.getColumnName(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(") ");
 
-        sb.append(" SELECT ");
-        for (int i=0; i < size; i++) {
-            if(i != 0) { sb.append(", "); }
-            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(" WHERE (SELECT result FROM upsert) is null");
+        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);
+        }
+    }
 }

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

@@ -1,40 +1,40 @@
-package org.embulk.output.postgresql;
-
-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 PostgreSQLOutputConnector
-        implements JdbcOutputConnector
-{
-    private static final Driver driver = new org.postgresql.Driver();
-
-    private final String url;
-    private final Properties properties;
-    private final String schemaName;
-
-    public PostgreSQLOutputConnector(String url, Properties properties, String schemaName)
-    {
-        this.url = url;
-        this.properties = properties;
-        this.schemaName = schemaName;
-    }
-
-    @Override
-    public PostgreSQLOutputConnection connect(boolean autoCommit) throws SQLException
-    {
-        Connection c = driver.connect(url, properties);
-        try {
-            PostgreSQLOutputConnection con = new PostgreSQLOutputConnection(c, schemaName, autoCommit);
-            c = null;
-            return con;
-        } finally {
-            if (c != null) {
-                c.close();
-            }
-        }
-    }
-}
+package org.embulk.output.postgresql;
+
+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 PostgreSQLOutputConnector
+        implements JdbcOutputConnector
+{
+    private static final Driver driver = new org.postgresql.Driver();
+
+    private final String url;
+    private final Properties properties;
+    private final String schemaName;
+
+    public PostgreSQLOutputConnector(String url, Properties properties, String schemaName)
+    {
+        this.url = url;
+        this.properties = properties;
+        this.schemaName = schemaName;
+    }
+
+    @Override
+    public PostgreSQLOutputConnection connect(boolean autoCommit) throws SQLException
+    {
+        Connection c = driver.connect(url, properties);
+        try {
+            PostgreSQLOutputConnection con = new PostgreSQLOutputConnection(c, schemaName, autoCommit);
+            c = null;
+            return con;
+        } finally {
+            if (c != null) {
+                c.close();
+            }
+        }
+    }
+}

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: embulk-output-postgresql
 version: !ruby/object:Gem::Version
-  version: 0.2.4
+  version: 0.3.0
 platform: ruby
 authors:
 - Sadayuki Furuhashi
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2015-05-12 00:00:00.000000000 Z
+date: 2015-05-19 00:00:00.000000000 Z
 dependencies: []
 description: Inserts or updates records to a table.
 email:
@@ -25,9 +25,8 @@ files:
 - src/main/java/org/embulk/output/postgresql/PostgreSQLCopyBatchInsert.java
 - src/main/java/org/embulk/output/postgresql/PostgreSQLOutputConnection.java
 - src/main/java/org/embulk/output/postgresql/PostgreSQLOutputConnector.java
-- src/main/java/org/embulk/output/postgresql/PostgresqlBatchUpsert.java
-- classpath/embulk-output-jdbc-0.2.4.jar
-- classpath/embulk-output-postgresql-0.2.4.jar
+- classpath/embulk-output-jdbc-0.3.0.jar
+- classpath/embulk-output-postgresql-0.3.0.jar
 - classpath/jna-4.1.0.jar
 - classpath/jna-platform-4.1.0.jar
 - classpath/postgresql-9.4-1200-jdbc41.jar

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

@@ -1,24 +0,0 @@
-package org.embulk.output.postgresql;
-
-import org.embulk.output.jdbc.JdbcOutputConnection;
-import org.embulk.output.jdbc.JdbcOutputConnector;
-import org.embulk.output.jdbc.JdbcSchema;
-import org.embulk.output.jdbc.StandardBatchInsert;
-
-import java.io.IOException;
-import java.sql.PreparedStatement;
-import java.sql.SQLException;
-
-public class PostgresqlBatchUpsert extends StandardBatchInsert {
-
-    public PostgresqlBatchUpsert(JdbcOutputConnector connector) throws IOException, SQLException {
-        super(connector);
-    }
-
-    protected PreparedStatement newPreparedStatement(JdbcOutputConnection connection,
-                                                     String loadTable, JdbcSchema insertSchema) throws SQLException
-    {
-        return connection.prepareUpsertSql(loadTable, insertSchema);
-    }
-
-}