embulk-output-postgresql 0.4.1 → 0.4.2

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,73 +1,71 @@
1
- package org.embulk.output.postgresql;
2
-
3
- import java.io.File;
4
- import java.io.IOException;
5
- import java.io.FileInputStream;
6
- import java.sql.Connection;
7
- import java.sql.SQLException;
8
- import org.slf4j.Logger;
9
- import org.postgresql.copy.CopyManager;
10
- import org.postgresql.core.BaseConnection;
11
- import org.embulk.spi.Exec;
12
- import org.embulk.output.jdbc.JdbcSchema;
13
-
14
- public class PostgreSQLCopyBatchInsert
15
- extends AbstractPostgreSQLCopyBatchInsert
16
- {
17
- private final Logger logger = Exec.getLogger(PostgreSQLCopyBatchInsert.class);
18
- private final PostgreSQLOutputConnector connector;
19
-
20
- private PostgreSQLOutputConnection connection = null;
21
- private CopyManager copyManager = null;
22
- private String copySql = null;
23
- private long totalRows;
24
-
25
- public PostgreSQLCopyBatchInsert(PostgreSQLOutputConnector connector) throws IOException, SQLException
26
- {
27
- super();
28
- this.connector = connector;
29
- }
30
-
31
- @Override
32
- public void prepare(String loadTable, JdbcSchema insertSchema) throws SQLException
33
- {
34
- this.connection = connector.connect(true);
35
- this.copySql = connection.buildCopySql(loadTable, insertSchema);
36
- this.copyManager = connection.newCopyManager();
37
- logger.info("Copy SQL: "+copySql);
38
- }
39
-
40
- @Override
41
- public void flush() throws IOException, SQLException
42
- {
43
- File file = closeCurrentFile(); // flush buffered data in writer
44
-
45
- logger.info(String.format("Loading %,d rows (%,d bytes)", batchRows, file.length()));
46
- long startTime = System.currentTimeMillis();
47
- FileInputStream in = new FileInputStream(file);
48
- try {
49
- // TODO check age of connection and call isValid if it's old and reconnect if it's invalid
50
- copyManager.copyIn(copySql, in);
51
- } finally {
52
- in.close();
53
- }
54
- double seconds = (System.currentTimeMillis() - startTime) / 1000.0;
55
-
56
- totalRows += batchRows;
57
- batchRows = 0;
58
- logger.info(String.format("> %.2f seconds (loaded %,d rows in total)", seconds, totalRows));
59
-
60
- openNewFile();
61
- file.delete();
62
- }
63
-
64
- @Override
65
- public void close() throws IOException, SQLException
66
- {
67
- closeCurrentFile().delete();
68
- if (connection != null) {
69
- connection.close();
70
- connection = null;
71
- }
72
- }
73
- }
1
+ package org.embulk.output.postgresql;
2
+
3
+ import java.io.File;
4
+ import java.io.IOException;
5
+ import java.io.FileInputStream;
6
+ import java.sql.SQLException;
7
+ import org.slf4j.Logger;
8
+ import org.postgresql.copy.CopyManager;
9
+ import org.embulk.spi.Exec;
10
+ import org.embulk.output.jdbc.JdbcSchema;
11
+
12
+ public class PostgreSQLCopyBatchInsert
13
+ extends AbstractPostgreSQLCopyBatchInsert
14
+ {
15
+ private final Logger logger = Exec.getLogger(PostgreSQLCopyBatchInsert.class);
16
+ private final PostgreSQLOutputConnector connector;
17
+
18
+ private PostgreSQLOutputConnection connection = null;
19
+ private CopyManager copyManager = null;
20
+ private String copySql = null;
21
+ private long totalRows;
22
+
23
+ public PostgreSQLCopyBatchInsert(PostgreSQLOutputConnector connector) throws IOException, SQLException
24
+ {
25
+ super();
26
+ this.connector = connector;
27
+ }
28
+
29
+ @Override
30
+ public void prepare(String loadTable, JdbcSchema insertSchema) throws SQLException
31
+ {
32
+ this.connection = connector.connect(true);
33
+ this.copySql = connection.buildCopySql(loadTable, insertSchema);
34
+ this.copyManager = connection.newCopyManager();
35
+ logger.info("Copy SQL: "+copySql);
36
+ }
37
+
38
+ @Override
39
+ public void flush() throws IOException, SQLException
40
+ {
41
+ File file = closeCurrentFile(); // flush buffered data in writer
42
+
43
+ logger.info(String.format("Loading %,d rows (%,d bytes)", batchRows, file.length()));
44
+ long startTime = System.currentTimeMillis();
45
+ FileInputStream in = new FileInputStream(file);
46
+ try {
47
+ // TODO check age of connection and call isValid if it's old and reconnect if it's invalid
48
+ copyManager.copyIn(copySql, in);
49
+ } finally {
50
+ in.close();
51
+ }
52
+ double seconds = (System.currentTimeMillis() - startTime) / 1000.0;
53
+
54
+ totalRows += batchRows;
55
+ batchRows = 0;
56
+ logger.info(String.format("> %.2f seconds (loaded %,d rows in total)", seconds, totalRows));
57
+
58
+ openNewFile();
59
+ file.delete();
60
+ }
61
+
62
+ @Override
63
+ public void close() throws IOException, SQLException
64
+ {
65
+ closeCurrentFile().delete();
66
+ if (connection != null) {
67
+ connection.close();
68
+ connection = null;
69
+ }
70
+ }
71
+ }
@@ -1,157 +1,156 @@
1
- package org.embulk.output.postgresql;
2
-
3
- import java.util.List;
4
- import java.sql.Connection;
5
- import java.sql.SQLException;
6
- import java.sql.Statement;
7
- import org.postgresql.copy.CopyManager;
8
- import org.postgresql.core.BaseConnection;
9
- import org.embulk.spi.Exec;
10
- import org.embulk.output.jdbc.JdbcOutputConnection;
11
- import org.embulk.output.jdbc.JdbcColumn;
12
- import org.embulk.output.jdbc.JdbcSchema;
13
-
14
- public class PostgreSQLOutputConnection
15
- extends JdbcOutputConnection
16
- {
17
- public PostgreSQLOutputConnection(Connection connection, String schemaName, boolean autoCommit)
18
- throws SQLException
19
- {
20
- super(connection, schemaName);
21
- connection.setAutoCommit(autoCommit);
22
- }
23
-
24
- public String buildCopySql(String toTable, JdbcSchema toTableSchema)
25
- {
26
- StringBuilder sb = new StringBuilder();
27
-
28
- sb.append("COPY ");
29
- quoteIdentifierString(sb, toTable);
30
- sb.append(" (");
31
- for (int i=0; i < toTableSchema.getCount(); i++) {
32
- if (i != 0) { sb.append(", "); }
33
- quoteIdentifierString(sb, toTableSchema.getColumnName(i));
34
- }
35
- sb.append(") ");
36
- sb.append("FROM STDIN");
37
-
38
- return sb.toString();
39
- }
40
-
41
- public CopyManager newCopyManager() throws SQLException
42
- {
43
- return new CopyManager((BaseConnection) connection);
44
- }
45
-
46
- @Override
47
- protected String buildCollectMergeSql(List<String> fromTables, JdbcSchema schema, String toTable, List<String> mergeKeys) throws SQLException
48
- {
49
- StringBuilder sb = new StringBuilder();
50
-
51
- sb.append("WITH updated AS (");
52
- sb.append("UPDATE ");
53
- quoteIdentifierString(sb, toTable);
54
- sb.append(" SET ");
55
- for (int i=0; i < schema.getCount(); i++) {
56
- if (i != 0) { sb.append(", "); }
57
- quoteIdentifierString(sb, schema.getColumnName(i));
58
- sb.append(" = S.");
59
- quoteIdentifierString(sb, schema.getColumnName(i));
60
- }
61
- sb.append(" FROM (");
62
- for (int i=0; i < fromTables.size(); i++) {
63
- if (i != 0) { sb.append(" UNION ALL "); }
64
- sb.append("SELECT ");
65
- for(int j=0; j < schema.getCount(); j++) {
66
- if (j != 0) { sb.append(", "); }
67
- quoteIdentifierString(sb, schema.getColumnName(j));
68
- }
69
- sb.append(" FROM ");
70
- quoteIdentifierString(sb, fromTables.get(i));
71
- }
72
- sb.append(") S");
73
- sb.append(" WHERE ");
74
- for (int i=0; i < mergeKeys.size(); i++) {
75
- if (i != 0) { sb.append(" AND "); }
76
- quoteIdentifierString(sb, toTable);
77
- sb.append(".");
78
- quoteIdentifierString(sb, mergeKeys.get(i));
79
- sb.append(" = ");
80
- sb.append("S.");
81
- quoteIdentifierString(sb, mergeKeys.get(i));
82
- }
83
- sb.append(" RETURNING ");
84
- for (int i=0; i < mergeKeys.size(); i++) {
85
- if (i != 0) { sb.append(", "); }
86
- sb.append("S.");
87
- quoteIdentifierString(sb, mergeKeys.get(i));
88
- }
89
- sb.append(") ");
90
-
91
- sb.append("INSERT INTO ");
92
- quoteIdentifierString(sb, toTable);
93
- sb.append(" (");
94
- for (int i=0; i < schema.getCount(); i++) {
95
- if (i != 0) { sb.append(", "); }
96
- quoteIdentifierString(sb, schema.getColumnName(i));
97
- }
98
- sb.append(") ");
99
- sb.append("SELECT DISTINCT ON (");
100
- for (int i=0; i < mergeKeys.size(); i++) {
101
- if (i != 0) { sb.append(", "); }
102
- quoteIdentifierString(sb, mergeKeys.get(i));
103
- }
104
- sb.append(") * FROM (");
105
- for (int i=0; i < fromTables.size(); i++) {
106
- if (i != 0) { sb.append(" UNION ALL "); }
107
- sb.append("SELECT ");
108
- for(int j=0; j < schema.getCount(); j++) {
109
- if (j != 0) { sb.append(", "); }
110
- quoteIdentifierString(sb, schema.getColumnName(j));
111
- }
112
- sb.append(" FROM ");
113
- quoteIdentifierString(sb, fromTables.get(i));
114
- }
115
- sb.append(") S ");
116
- sb.append("WHERE NOT EXISTS (");
117
- sb.append("SELECT 1 FROM updated WHERE ");
118
- for (int i=0; i < mergeKeys.size(); i++) {
119
- if (i != 0) { sb.append(" AND "); }
120
- sb.append("S.");
121
- quoteIdentifierString(sb, mergeKeys.get(i));
122
- sb.append(" = ");
123
- sb.append("updated.");
124
- quoteIdentifierString(sb, mergeKeys.get(i));
125
- }
126
- sb.append(") ");
127
-
128
- return sb.toString();
129
- }
130
-
131
- protected void collectReplaceView(List<String> fromTables, JdbcSchema schema, String toTable) throws SQLException
132
- {
133
- Statement stmt = connection.createStatement();
134
- try {
135
- String sql = buildCollectInsertSql(fromTables, schema, toTable);
136
- executeUpdate(stmt, sql);
137
- commitIfNecessary(connection);
138
- } catch (SQLException ex) {
139
- throw safeRollback(connection, ex);
140
- } finally {
141
- stmt.close();
142
- }
143
- }
144
-
145
- @Override
146
- protected String buildColumnTypeName(JdbcColumn c)
147
- {
148
- switch(c.getSimpleTypeName()) {
149
- case "CLOB":
150
- return "TEXT";
151
- case "BLOB":
152
- return "BYTEA";
153
- default:
154
- return super.buildColumnTypeName(c);
155
- }
156
- }
157
- }
1
+ package org.embulk.output.postgresql;
2
+
3
+ import java.util.List;
4
+ import java.sql.Connection;
5
+ import java.sql.SQLException;
6
+ import java.sql.Statement;
7
+ import org.postgresql.copy.CopyManager;
8
+ import org.postgresql.core.BaseConnection;
9
+ import org.embulk.output.jdbc.JdbcOutputConnection;
10
+ import org.embulk.output.jdbc.JdbcColumn;
11
+ import org.embulk.output.jdbc.JdbcSchema;
12
+
13
+ public class PostgreSQLOutputConnection
14
+ extends JdbcOutputConnection
15
+ {
16
+ public PostgreSQLOutputConnection(Connection connection, String schemaName, boolean autoCommit)
17
+ throws SQLException
18
+ {
19
+ super(connection, schemaName);
20
+ connection.setAutoCommit(autoCommit);
21
+ }
22
+
23
+ public String buildCopySql(String toTable, JdbcSchema toTableSchema)
24
+ {
25
+ StringBuilder sb = new StringBuilder();
26
+
27
+ sb.append("COPY ");
28
+ quoteIdentifierString(sb, toTable);
29
+ sb.append(" (");
30
+ for (int i=0; i < toTableSchema.getCount(); i++) {
31
+ if (i != 0) { sb.append(", "); }
32
+ quoteIdentifierString(sb, toTableSchema.getColumnName(i));
33
+ }
34
+ sb.append(") ");
35
+ sb.append("FROM STDIN");
36
+
37
+ return sb.toString();
38
+ }
39
+
40
+ public CopyManager newCopyManager() throws SQLException
41
+ {
42
+ return new CopyManager((BaseConnection) connection);
43
+ }
44
+
45
+ @Override
46
+ protected String buildCollectMergeSql(List<String> fromTables, JdbcSchema schema, String toTable, List<String> mergeKeys) throws SQLException
47
+ {
48
+ StringBuilder sb = new StringBuilder();
49
+
50
+ sb.append("WITH updated AS (");
51
+ sb.append("UPDATE ");
52
+ quoteIdentifierString(sb, toTable);
53
+ sb.append(" SET ");
54
+ for (int i=0; i < schema.getCount(); i++) {
55
+ if (i != 0) { sb.append(", "); }
56
+ quoteIdentifierString(sb, schema.getColumnName(i));
57
+ sb.append(" = S.");
58
+ quoteIdentifierString(sb, schema.getColumnName(i));
59
+ }
60
+ sb.append(" FROM (");
61
+ for (int i=0; i < fromTables.size(); i++) {
62
+ if (i != 0) { sb.append(" UNION ALL "); }
63
+ sb.append("SELECT ");
64
+ for(int j=0; j < schema.getCount(); j++) {
65
+ if (j != 0) { sb.append(", "); }
66
+ quoteIdentifierString(sb, schema.getColumnName(j));
67
+ }
68
+ sb.append(" FROM ");
69
+ quoteIdentifierString(sb, fromTables.get(i));
70
+ }
71
+ sb.append(") S");
72
+ sb.append(" WHERE ");
73
+ for (int i=0; i < mergeKeys.size(); i++) {
74
+ if (i != 0) { sb.append(" AND "); }
75
+ quoteIdentifierString(sb, toTable);
76
+ sb.append(".");
77
+ quoteIdentifierString(sb, mergeKeys.get(i));
78
+ sb.append(" = ");
79
+ sb.append("S.");
80
+ quoteIdentifierString(sb, mergeKeys.get(i));
81
+ }
82
+ sb.append(" RETURNING ");
83
+ for (int i=0; i < mergeKeys.size(); i++) {
84
+ if (i != 0) { sb.append(", "); }
85
+ sb.append("S.");
86
+ quoteIdentifierString(sb, mergeKeys.get(i));
87
+ }
88
+ sb.append(") ");
89
+
90
+ sb.append("INSERT INTO ");
91
+ quoteIdentifierString(sb, toTable);
92
+ sb.append(" (");
93
+ for (int i=0; i < schema.getCount(); i++) {
94
+ if (i != 0) { sb.append(", "); }
95
+ quoteIdentifierString(sb, schema.getColumnName(i));
96
+ }
97
+ sb.append(") ");
98
+ sb.append("SELECT DISTINCT ON (");
99
+ for (int i=0; i < mergeKeys.size(); i++) {
100
+ if (i != 0) { sb.append(", "); }
101
+ quoteIdentifierString(sb, mergeKeys.get(i));
102
+ }
103
+ sb.append(") * FROM (");
104
+ for (int i=0; i < fromTables.size(); i++) {
105
+ if (i != 0) { sb.append(" UNION ALL "); }
106
+ sb.append("SELECT ");
107
+ for(int j=0; j < schema.getCount(); j++) {
108
+ if (j != 0) { sb.append(", "); }
109
+ quoteIdentifierString(sb, schema.getColumnName(j));
110
+ }
111
+ sb.append(" FROM ");
112
+ quoteIdentifierString(sb, fromTables.get(i));
113
+ }
114
+ sb.append(") S ");
115
+ sb.append("WHERE NOT EXISTS (");
116
+ sb.append("SELECT 1 FROM updated WHERE ");
117
+ for (int i=0; i < mergeKeys.size(); i++) {
118
+ if (i != 0) { sb.append(" AND "); }
119
+ sb.append("S.");
120
+ quoteIdentifierString(sb, mergeKeys.get(i));
121
+ sb.append(" = ");
122
+ sb.append("updated.");
123
+ quoteIdentifierString(sb, mergeKeys.get(i));
124
+ }
125
+ sb.append(") ");
126
+
127
+ return sb.toString();
128
+ }
129
+
130
+ protected void collectReplaceView(List<String> fromTables, JdbcSchema schema, String toTable) throws SQLException
131
+ {
132
+ Statement stmt = connection.createStatement();
133
+ try {
134
+ String sql = buildCollectInsertSql(fromTables, schema, toTable);
135
+ executeUpdate(stmt, sql);
136
+ commitIfNecessary(connection);
137
+ } catch (SQLException ex) {
138
+ throw safeRollback(connection, ex);
139
+ } finally {
140
+ stmt.close();
141
+ }
142
+ }
143
+
144
+ @Override
145
+ protected String buildColumnTypeName(JdbcColumn c)
146
+ {
147
+ switch(c.getSimpleTypeName()) {
148
+ case "CLOB":
149
+ return "TEXT";
150
+ case "BLOB":
151
+ return "BYTEA";
152
+ default:
153
+ return super.buildColumnTypeName(c);
154
+ }
155
+ }
156
+ }