Merge pull request #220 from embulk/add_statement_timeout_to_postgre

Add statement timeout for PostgreSQL

Author: dmikurube

embulk-input-postgresql/src/main/java/org/embulk/input/PostgreSQLInputPlugin.java

@@ -53,6 +53,10 @@ public interface PostgreSQLPluginTask
         @Config("application_name")
         @ConfigDefault("\"embulk-input-postgresql\"")
         public String getApplicationName();
+
+        @Config("statement_timeout_millis")
+        @ConfigDefault("null")
+        public Optional<Integer> getStatementTimeoutMillis();
     }
 
     @Override
@@ -96,7 +100,7 @@ protected PostgreSQLInputConnection newConnection(PluginTask task) throws SQLExc
 
         Connection con = DriverManager.getConnection(url, props);
         try {
-            PostgreSQLInputConnection c = new PostgreSQLInputConnection(con, t.getSchema());
+            PostgreSQLInputConnection c = new PostgreSQLInputConnection(con, t.getSchema(), t.getStatementTimeoutMillis());
             con = null;
             return c;
         } finally {

embulk-input-postgresql/src/main/java/org/embulk/input/postgresql/PostgreSQLInputConnection.java

@@ -1,10 +1,13 @@
 package org.embulk.input.postgresql;
 
+import java.sql.Statement;
 import java.util.List;
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.util.Optional;
+
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.embulk.input.jdbc.JdbcInputConnection;
@@ -16,10 +19,11 @@ public class PostgreSQLInputConnection
 {
     private static final Logger logger = LoggerFactory.getLogger(PostgreSQLInputConnection.class);
 
-    public PostgreSQLInputConnection(Connection connection, String schemaName)
+    public PostgreSQLInputConnection(Connection connection, String schemaName, Optional<Integer> statementTimeoutMillis)
             throws SQLException
     {
         super(connection, schemaName);
+        setStatementTimeoutIfSpecified(statementTimeoutMillis);
     }
 
     @Override
@@ -76,4 +80,19 @@ public void close() throws SQLException
             // TODO close?
         }
     }
+
+    private void setStatementTimeoutIfSpecified(Optional<Integer> statementTimeoutMillis)
+        throws SQLException
+    {
+        if (statementTimeoutMillis.isPresent() && statementTimeoutMillis.get() > 0) {
+            Statement stmt = connection.createStatement();
+            try {
+                String sql = "SET statement_timeout TO " + quoteIdentifierString(String.valueOf(statementTimeoutMillis.get()));
+                executeUpdate(sql);
+            }
+            finally {
+                stmt.close();
+            }
+        }
+    }
 }

embulk-input-redshift/src/main/java/org/embulk/input/RedshiftInputPlugin.java

@@ -1,5 +1,6 @@
 package org.embulk.input;
 
+import java.util.Optional;
 import java.util.Properties;
 import java.sql.Connection;
 import java.sql.Driver;
@@ -46,6 +47,10 @@ public interface RedshiftPluginTask
         @Config("ssl")
         @ConfigDefault("false")
         public boolean getSsl();
+
+        @Config("statement_timeout_millis")
+        @ConfigDefault("null")
+        public Optional<Integer> getStatementTimeoutMillis();
     }
 
     @Override
@@ -85,7 +90,7 @@ protected PostgreSQLInputConnection newConnection(PluginTask task) throws SQLExc
 
         Connection con = driver.connect(url, props);
         try {
-            PostgreSQLInputConnection c = new PostgreSQLInputConnection(con, t.getSchema());
+            PostgreSQLInputConnection c = new PostgreSQLInputConnection(con, t.getSchema(), t.getStatementTimeoutMillis());
             con = null;
             return c;
         } finally {