Adds support for timeouts in select queries (#123) (#124)

* Adds support for timeouts in select queries (#123)

This provides a mechanism to specify a default timeout for all select queries created by a given
DatabaseEngine plus the possibility to override the configured value on individual queries. It
also exposes the Statement#cancel() method, that allows a thread to cancel a query running on other
query and is supported by the production-ready JDBC drivers supported by pdb.

src/main/java/com/feedzai/commons/sql/abstraction/dml/result/MySqlResultIterator.java

@@ -17,6 +17,7 @@
 
 import com.feedzai.commons.sql.abstraction.engine.DatabaseEngineException;
 import com.feedzai.commons.sql.abstraction.engine.impl.MySqlEngine;
+import com.mysql.jdbc.exceptions.MySQLTimeoutException;
 
 import java.sql.PreparedStatement;
 import java.sql.Statement;
@@ -53,4 +54,9 @@ public MySqlResultIterator(PreparedStatement statement) throws DatabaseEngineExc
     public ResultColumn createResultColumn(String name, Object value) {
         return new MySqlResultColumn(name, value);
     }
+
+    @Override
+    protected boolean isTimeoutException(final Exception exception) {
+        return exception instanceof MySQLTimeoutException;
+    }
 }

src/main/java/com/feedzai/commons/sql/abstraction/dml/result/PostgreSqlResultIterator.java

@@ -19,6 +19,7 @@
 import com.feedzai.commons.sql.abstraction.engine.impl.PostgreSqlEngine;
 
 import java.sql.PreparedStatement;
+import java.sql.SQLException;
 import java.sql.Statement;
 
 /**
@@ -28,6 +29,12 @@
  * @since 2.0.0
  */
 public class PostgreSqlResultIterator extends ResultIterator {
+
+    /**
+     * The SQL State that indicates a timeout occurred.
+     */
+    private static final String TIMEOUT_SQL_STATE = "57014";
+
     /**
      * Creates a new instance of {@link PostgreSqlResultIterator}.
      *
@@ -53,4 +60,10 @@ public PostgreSqlResultIterator(PreparedStatement statement) throws DatabaseEngi
     public ResultColumn createResultColumn(String name, Object value) {
         return new PostgreSqlResultColumn(name, value);
     }
+
+    @Override
+    protected boolean isTimeoutException(Exception exception) {
+        return super.isTimeoutException (exception) ||
+                exception instanceof SQLException && TIMEOUT_SQL_STATE.equals(((SQLException) exception).getSQLState());
+    }
 }

src/main/java/com/feedzai/commons/sql/abstraction/dml/result/ResultIterator.java

@@ -16,12 +16,15 @@
 package com.feedzai.commons.sql.abstraction.dml.result;
 
 import com.feedzai.commons.sql.abstraction.engine.DatabaseEngineException;
+import com.feedzai.commons.sql.abstraction.engine.DatabaseEngineTimeoutException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.ResultSetMetaData;
+import java.sql.SQLException;
+import java.sql.SQLTimeoutException;
 import java.sql.Statement;
 import java.util.ArrayList;
 import java.util.LinkedHashMap;
@@ -40,6 +43,7 @@ public abstract class ResultIterator implements AutoCloseable {
      * The logger.
      */
     private static final Logger logger = LoggerFactory.getLogger(ResultIterator.class);
+
     /**
      * The statement.
      */
@@ -98,10 +102,23 @@ private ResultIterator(Statement statement, String sql, boolean isPreparedStatem
 
         } catch (final Exception e) {
             close();
-            throw new DatabaseEngineException("Could not process result set.", e);
+            throw (isTimeoutException(e) ?
+                new DatabaseEngineTimeoutException("Timeout waiting for query execution", e) :
+                new DatabaseEngineException("Could not process result set.", e));
         }
     }
 
+    /**
+     * Indicates if a given exception is a timeout. Logic for this may be driver-specific, so
+     * drivers that support query timeouts may have to override this method.
+     *
+     * @param exception  The exception to check.
+     * @return {@code true} if the exception is a timeout, {@code false} otherwise.
+     */
+    protected boolean isTimeoutException(final Exception exception) {
+        return (exception instanceof SQLTimeoutException);
+    }
+
     /**
      * Creates a new instance of {@link ResultIterator} for regular {@link Statement}.
      *
@@ -225,6 +242,31 @@ public List<String> getColumnNames() {
         return columnNames;
     }
 
+    /**
+     * Attempts to cancel the current query. This relies on the JDBC driver supporting
+     * {@link Statement#cancel()}, which is not guaranteed on all drivers.
+     *
+     * A possible use case for this method is to implement a timeout; If that's the case, see also
+     * {@link com.feedzai.commons.sql.abstraction.engine.AbstractDatabaseEngine#iterator(String, int, int)} for
+     * an alternative way to accomplish this.
+     *
+     * This method is expected to be invoked from a thread distinct of the one that is reading
+     * from the result set.
+     *
+     * @return {@code true} if the query was cancelled, {@code false} otherwise.
+     */
+    public boolean cancel() {
+        try {
+            if (!closed) {
+                statement.cancel();
+            }
+            return true;
+        } catch (SQLException ex) {
+            logger.debug("Could not cancel statement", ex);
+            return false;
+        }
+    }
+
     /**
      * Closes the {@link ResultSet} and the {@link Statement} if applicable.
      */

src/main/java/com/feedzai/commons/sql/abstraction/engine/AbstractDatabaseEngine.java

@@ -33,6 +33,7 @@
 import com.feedzai.commons.sql.abstraction.engine.handler.OperationFault;
 import com.feedzai.commons.sql.abstraction.entry.EntityEntry;
 import com.feedzai.commons.sql.abstraction.util.AESHelper;
+import com.feedzai.commons.sql.abstraction.util.Constants;
 import com.feedzai.commons.sql.abstraction.util.InitiallyReusableByteArrayOutputStream;
 import com.feedzai.commons.sql.abstraction.util.PreparedStatementCapsule;
 import com.google.common.collect.ImmutableList;
@@ -75,6 +76,7 @@
 import static com.feedzai.commons.sql.abstraction.engine.configuration.PdbProperties.ENCRYPTED_USERNAME;
 import static com.feedzai.commons.sql.abstraction.engine.configuration.PdbProperties.JDBC;
 import static com.feedzai.commons.sql.abstraction.engine.configuration.PdbProperties.SECRET_LOCATION;
+import static com.feedzai.commons.sql.abstraction.util.Constants.NO_TIMEOUT;
 import static com.feedzai.commons.sql.abstraction.util.StringUtils.quotize;
 import static com.feedzai.commons.sql.abstraction.util.StringUtils.readString;
 
@@ -860,6 +862,22 @@ public synchronized int executeUpdate(final String query) throws DatabaseEngineE
         }
     }
 
+    /**
+     * Creates a {@link Statement} that will be used for selects, i.e., may have an associated
+     * read timeout.
+     *
+     * @param readTimeout   The timeout.
+     * @return The {@link Statement}
+     * @throws SQLException If there is an error creating the statement.
+     */
+    protected Statement createSelectStatement(int readTimeout) throws SQLException {
+        final Statement s = conn.createStatement();
+        if (readTimeout != NO_TIMEOUT) {
+            s.setQueryTimeout(readTimeout);
+        }
+        return s;
+    }
+
     /**
      * Executes the given update.
      *
@@ -1106,29 +1124,39 @@ protected int entityToPreparedStatementForBatch(final DbEntity entity, final Pre
      * @throws DatabaseEngineException If something goes wrong executing the query.
      */
     @Override
-    public synchronized List<Map<String, ResultColumn>> query(final Expression query) throws DatabaseEngineException {
+    public List<Map<String, ResultColumn>> query(final Expression query) throws DatabaseEngineException {
         return query(translate(query));
     }
 
+    @Override
+    public List<Map<String, ResultColumn>> query(Expression query, int readTimeoutOverride) throws DatabaseEngineException {
+        return query(translate(query), readTimeoutOverride);
+    }
+
     /**
      * Executes the given query.
      *
      * @param query The query to execute.
      * @throws DatabaseEngineException If something goes wrong executing the query.
      */
     @Override
-    public synchronized List<Map<String, ResultColumn>> query(final String query) throws DatabaseEngineException {
+    public List<Map<String, ResultColumn>> query(final String query) throws DatabaseEngineException {
         return processResultIterator(iterator(query));
     }
 
+    @Override
+    public List<Map<String, ResultColumn>> query(String query, int readTimeoutOverride) throws DatabaseEngineException {
+        return processResultIterator(iterator(query, properties.getFetchSize(), readTimeoutOverride));
+    }
+
     /**
      * Process a whole {@link ResultIterator}.
      *
      * @param it The iterator.
      * @return A list of rows in the form of {@link Map}.
      * @throws DatabaseEngineException If a database access error occurs.
      */
-    protected List<Map<String, ResultColumn>> processResultIterator(ResultIterator it) throws DatabaseEngineException {
+    protected synchronized List<Map<String, ResultColumn>> processResultIterator(ResultIterator it) throws DatabaseEngineException {
         List<Map<String, ResultColumn>> res = new ArrayList<>();
 
         Map<String, ResultColumn> temp;
@@ -1140,34 +1168,47 @@ protected List<Map<String, ResultColumn>> processResultIterator(ResultIterator i
     }
 
     @Override
-    public synchronized ResultIterator iterator(String query) throws DatabaseEngineException {
+    public ResultIterator iterator(String query) throws DatabaseEngineException {
         return iterator(query, properties.getFetchSize());
     }
 
     @Override
-    public synchronized ResultIterator iterator(Expression query) throws DatabaseEngineException {
+    public ResultIterator iterator(Expression query) throws DatabaseEngineException {
         return iterator(query, properties.getFetchSize());
     }
 
     @Override
     public ResultIterator iterator(String query, int fetchSize) throws DatabaseEngineException {
+        return iterator(query, fetchSize, properties.getSelectQueryTimeout());
+    }
+
+    @Override
+    public ResultIterator iterator(Expression query, int fetchSize) throws DatabaseEngineException {
+        return iterator(translate(query), fetchSize);
+    }
+
+    @Override
+    public ResultIterator iterator(Expression query, int fetchSize, int readTimeoutOverride) throws DatabaseEngineException {
+        return iterator(translate(query), fetchSize, readTimeoutOverride);
+    }
+
+    @Override
+    public synchronized ResultIterator iterator(String query, int fetchSize, int readTimeoutOverride) throws DatabaseEngineException {
         try {
             getConnection();
-            Statement stmt = conn.createStatement();
+            Statement stmt = createSelectStatement(readTimeoutOverride);
             stmt.setFetchSize(fetchSize);
             logger.trace(query);
             return createResultIterator(stmt, query);
 
+        } catch (final DatabaseEngineTimeoutException e) {
+            throw e;
+
         } catch (final Exception e) {
             throw new DatabaseEngineException("Error querying", e);
         }
     }
 
-    @Override
-    public ResultIterator iterator(Expression query, int fetchSize) throws DatabaseEngineException {
-        return iterator(translate(query), fetchSize);
-    }
-
     /**
      * Creates a specific {@link ResultIterator} given the engine implementation.
      *
@@ -1484,7 +1525,7 @@ public Map<String, DbColumnType> getQueryMetadata(String query) throws DatabaseE
 
         try {
             getConnection();
-            stmt = conn.createStatement();
+            stmt = createSelectStatement(Constants.NO_TIMEOUT);  // No timeout on metadata queries
             long start = System.currentTimeMillis();
             rs = stmt.executeQuery(query);
             logger.trace("[{} ms] {}", (System.currentTimeMillis() - start), query);

src/main/java/com/feedzai/commons/sql/abstraction/engine/DatabaseEngine.java

@@ -303,6 +303,15 @@ public interface DatabaseEngine extends AutoCloseable {
      */
     List<Map<String, ResultColumn>> query(final Expression query) throws DatabaseEngineException;
 
+    /**
+     * Executes the given query overriding the configured query timeout (see {@link PdbProperties#getSelectQueryTimeout()}).
+     *
+     * @param query The query to execute.
+     * @param readTimeoutOverride The query timeout to use.
+     * @throws DatabaseEngineException If something goes wrong executing the query.
+     */
+    List<Map<String, ResultColumn>> query(final Expression query, final int readTimeoutOverride) throws DatabaseEngineException;
+
     /**
      * Executes the given native query.
      *
@@ -311,6 +320,15 @@ public interface DatabaseEngine extends AutoCloseable {
      */
     List<Map<String, ResultColumn>> query(final String query) throws DatabaseEngineException;
 
+    /**
+     * Executes the given native query overriding the configured query timeout (see {@link PdbProperties#getSelectQueryTimeout()}).
+     *
+     * @param query The query to execute.
+     * @param readTimeoutOverride The query timeout to use.
+     * @throws DatabaseEngineException If something goes wrong executing the query.
+     */
+    List<Map<String, ResultColumn>> query(final String query, final int readTimeoutOverride) throws DatabaseEngineException;
+
     /**
      * Gets the database entities for the current schema.
      *
@@ -572,6 +590,18 @@ void createPreparedStatement(final String name, final String query, final int ti
      */
     ResultIterator iterator(final String query, final int fetchSize) throws DatabaseEngineException;
 
+    /**
+     * Creates an iterator for the given SQL sentence overriding the configured query
+     * timeout (see {@link PdbProperties#getSelectQueryTimeout()}).
+     *
+     * @param query     The query.
+     * @param fetchSize The number of rows to fetch each time.
+     * @param readTimeoutOverride The query timeout to use.
+     * @return An iterator for the results of the given SQL query.
+     * @throws DatabaseEngineException If a database access error occurs.
+     */
+    ResultIterator iterator(final String query, final int fetchSize, final int readTimeoutOverride) throws DatabaseEngineException;
+
     /**
      * Creates an iterator for the given SQL expression.
      *
@@ -582,6 +612,18 @@ void createPreparedStatement(final String name, final String query, final int ti
      */
     ResultIterator iterator(final Expression query, final int fetchSize) throws DatabaseEngineException;
 
+    /**
+     * Creates an iterator for the given SQL expression overriding the configured query
+     * timeout (see {@link PdbProperties#getSelectQueryTimeout()}).
+     *
+     * @param query     The expression that represents the query.
+     * @param fetchSize The number of rows to fetch each time.
+     * @param readTimeoutOverride The query timeout to use.
+     * @return An iterator for the results of the given SQL expression.
+     * @throws DatabaseEngineException If a database access error occurs.
+     */
+    ResultIterator iterator(final Expression query, final int fetchSize, final int readTimeoutOverride) throws DatabaseEngineException;
+
     /**
      * Creates an iterator for the {@link java.sql.PreparedStatement} bound to the given name.
      *

src/main/java/com/feedzai/commons/sql/abstraction/engine/DatabaseEngineTimeoutException.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2019 Feedzai
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.feedzai.commons.sql.abstraction.engine;
+
+/**
+ * A {@link DatabaseEngineException} that represents a timeout error.
+ */
+public class DatabaseEngineTimeoutException extends DatabaseEngineException {
+
+    /**
+     * Constructs a new exception with the specified detail message and
+     * cause.  <p>Note that the detail message associated with
+     * {@code cause} is <i>not</i> automatically incorporated in
+     * this exception's detail message.
+     *
+     * @param message the detail message (which is saved for later retrieval
+     *                by the {@link #getMessage()} method).
+     * @param cause   the cause (which is saved for later retrieval by the
+     *                {@link #getCause()} method).  (A <tt>null</tt> value is
+     *                permitted, and indicates that the cause is nonexistent or
+     *                unknown.)
+     */
+    public DatabaseEngineTimeoutException(final String message, final Throwable cause) {
+        super(message, cause);
+    }
+}