embulk-output-postgresql 0.2.2 → 0.2.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: 45b9b9f967c210d54dc95836f1bedba07e37f879
-  data.tar.gz: 4228760b850eb3a3b9bca1af22d223f96967b628
+  metadata.gz: a9e12ceb372cd836d2bda55df13f1bd78ca1400a
+  data.tar.gz: 95cfefc30482d521bd3bcef520e494b7e5d82d39
 SHA512:
-  metadata.gz: ce4a7b37a25702f02dd927c58350d06282c3e13a0442093bbdf2b96bf024da9f6deea57fef9a35b1d35a9d27aa9e2867ae78c939a5370652db97bb67e9a63c06
-  data.tar.gz: 72fa53a7fd9e18954ba5da5e3d01fb3f266e395da8df0f274a7a9a49c8ff29ebc5b51631666b59be39f12a632a57856c2afa4bc566080a996324f5f58617aba9
+  metadata.gz: 29dde2105f5caf3e90cf6b423513403c82f87d031cf93df01320ed4ca0d4bec04edbcdbf6441bf9b95e9f5e274d1b8b3d32bd393f05aca02622d04a8b92222ef
+  data.tar.gz: 0a6ef290f497fe919486dac90dfc87e0eac002914deca3dd02f0f58ec10a7d10ff7e998d81f2d09794c06363abd100137f28e2be2202b6c5fffb4408be049340

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

@@ -1,8 +1,12 @@
 package org.embulk.output;
 
+import java.util.List;
 import java.util.Properties;
 import java.io.IOException;
 import java.sql.SQLException;
+
+import org.embulk.output.jdbc.setter.ColumnSetter;
+import org.embulk.output.postgresql.PostgresqlBatchUpsert;
 import org.embulk.spi.Exec;
 import org.embulk.config.Config;
 import org.embulk.config.ConfigDefault;
@@ -10,6 +14,7 @@ import org.embulk.output.jdbc.AbstractJdbcOutputPlugin;
 import org.embulk.output.jdbc.BatchInsert;
 import org.embulk.output.postgresql.PostgreSQLOutputConnector;
 import org.embulk.output.postgresql.PostgreSQLCopyBatchInsert;
+import org.embulk.spi.PageReader;
 
 public class PostgreSQLOutputPlugin
         extends AbstractJdbcOutputPlugin
@@ -85,9 +90,49 @@ public class PostgreSQLOutputPlugin
         return new PostgreSQLOutputConnector(url, props, t.getSchema());
     }
 
+    @Override
+    protected PluginPageOutput newPluginPageOutput(PageReader reader,
+                                                   BatchInsert batch, List<ColumnSetter> columnSetters,
+                                                   int batchSize)
+    {
+        return new PostgresPluginPageOutput(reader, batch, columnSetters, batchSize);
+    }
+
+    public static class PostgresPluginPageOutput extends PluginPageOutput
+    {
+
+        public PostgresPluginPageOutput(PageReader pageReader, BatchInsert batch, List<ColumnSetter> columnSetters, int batchSize) {
+            super(pageReader, batch, columnSetters, batchSize);
+        }
+
+        @Override
+        protected void handleColumnsSetters()
+        {
+            int size = columnSetters.size();
+            for (int i=0; i < size; i++) {
+                ColumnSetter columnSetter = columnSetters.get(i);
+                if (!columnSetter.getColumn().isPrimaryKey()) {
+                    columns.get(i).visit(columnSetter);
+                }
+            }
+            for (int i=0; i < size; i++) {
+                ColumnSetter columnSetter = columnSetters.get(i);
+                if (columnSetter.getColumn().isPrimaryKey()) {
+                    columns.get(i).visit(columnSetter);
+                }
+            }
+            for (int i=0; i < size; i++) {
+                columns.get(i).visit(columnSetters.get(i));
+            }
+        }
+
+    }
+
     @Override
     protected BatchInsert newBatchInsert(PluginTask task) throws IOException, SQLException
     {
-        return new PostgreSQLCopyBatchInsert(getConnector(task, true));
+        PostgreSQLOutputConnector connector = getConnector(task, true);
+        return task.getMode().isMerge() ? new PostgresqlBatchUpsert(connector) :
+                new PostgreSQLCopyBatchInsert(getConnector(task, true));
     }
 }

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

@@ -54,4 +54,55 @@ public class PostgreSQLOutputConnection
             return typeName;
         }
     }
+
+    @Override
+    protected String buildPrepareUpsertSql(String toTable, JdbcSchema toTableSchema) 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(" 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("=?");
+            }
+        }
+        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(")");
+
+        sb.append(" SELECT ");
+        for (int i=0; i < size; i++) {
+            if(i != 0) { sb.append(", "); }
+            sb.append("?");
+        }
+        sb.append(" WHERE (SELECT result FROM upsert) is null");
+
+        return sb.toString();
+    }
+
 }

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

@@ -0,0 +1,24 @@
+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);
+    }
+
+}

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: embulk-output-postgresql
 version: !ruby/object:Gem::Version
-  version: 0.2.2
+  version: 0.2.3
 platform: ruby
 authors:
 - Sadayuki Furuhashi
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2015-04-07 00:00:00.000000000 Z
+date: 2015-04-23 00:00:00.000000000 Z
 dependencies: []
 description: Inserts or updates records to a table.
 email:
@@ -25,12 +25,13 @@ 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
-- classpath/embulk-output-jdbc-0.2.2.jar
-- classpath/embulk-output-postgresql-0.2.2.jar
-- classpath/jna-4.1.0.jar
-- classpath/jna-platform-4.1.0.jar
+- src/main/java/org/embulk/output/postgresql/PostgresqlBatchUpsert.java
+- classpath/embulk-output-jdbc-0.2.3.jar
 - classpath/postgresql-9.4-1200-jdbc41.jar
+- classpath/embulk-output-postgresql-0.2.3.jar
 - classpath/slf4j-simple-1.7.7.jar
+- classpath/jna-platform-4.1.0.jar
+- classpath/jna-4.1.0.jar
 - classpath/waffle-jna-1.7.jar
 homepage: https://github.com/embulk/embulk-output-jdbc
 licenses: