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.
feedzai · Nov 12, 2019 · 8c2dc7a · 8c2dc7a
1 parent e5c3a55
commit 8c2dc7a
Show file tree

Hide file tree

Showing 6 changed files with 183 additions and 8 deletions.
diff --git a/src/main/java/com/feedzai/commons/sql/abstraction/dml/result/ResultIterator.java b/src/main/java/com/feedzai/commons/sql/abstraction/dml/result/ResultIterator.java
@@ -16,12 +16,14 @@
 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.Statement;
 import java.util.ArrayList;
 import java.util.LinkedHashMap;
@@ -40,6 +42,11 @@ public abstract class ResultIterator implements AutoCloseable {
      * The logger.
      */
     private static final Logger logger = LoggerFactory.getLogger(ResultIterator.class);
+
+    /**
+     * The SQL State that indicates a timeout occurred.
+     */
+    private static final String TIMEOUT_SQL_STATE = "57014";
     /**
      * The statement.
      */
@@ -98,10 +105,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. The default checks for SQL state 57014 (query
+     * cancelled by user), this method should be overrided for drivers where this is not the case.
+     *
+     * @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 SQLException && TIMEOUT_SQL_STATE.equals(((SQLException) exception).getSQLState()));
+    }
+
     /**
      * Creates a new instance of {@link ResultIterator} for regular {@link Statement}.
      *
@@ -225,6 +245,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.
      */

diff --git a/src/main/java/com/feedzai/commons/sql/abstraction/engine/AbstractDatabaseEngine.java b/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_SELECT_TIMEOUT;
 import static com.feedzai.commons.sql.abstraction.util.StringUtils.quotize;
 import static com.feedzai.commons.sql.abstraction.util.StringUtils.readString;
 
@@ -863,6 +865,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_SELECT_TIMEOUT) {
+            s.setQueryTimeout(readTimeout);
+        }
+        return s;
+    }
+
     /**
      * Executes the given update.
      *
@@ -1113,6 +1131,11 @@ public synchronized List<Map<String, ResultColumn>> query(final Expression query
         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.
      *
@@ -1124,6 +1147,11 @@ public synchronized List<Map<String, ResultColumn>> query(final String query) th
         return processResultIterator(iterator(query));
     }
 
+    @Override
+    public List<Map<String, ResultColumn>> query(String query, int readTimeoutOverride) throws DatabaseEngineException {
+        return processResultIterator(iterator(query, DEFAULT_FETCH_SIZE, readTimeoutOverride));
+    }
+
     /**
      * Process a whole {@link ResultIterator}.
      *
@@ -1154,9 +1182,19 @@ public synchronized ResultIterator iterator(Expression query) throws DatabaseEng
 
     @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(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);
@@ -1166,11 +1204,6 @@ public ResultIterator iterator(String query, int fetchSize) throws DatabaseEngin
         }
     }
 
-    @Override
-    public ResultIterator iterator(Expression query, int fetchSize) throws DatabaseEngineException {
-        return iterator(translate(query), fetchSize);
-    }
-
     /**
      * Creates a specific {@link ResultIterator} given the engine implementation.
      *
@@ -1487,7 +1520,7 @@ public Map<String, DbColumnType> getQueryMetadata(String query) throws DatabaseE
 
         try {
             getConnection();
-            stmt = conn.createStatement();
+            stmt = createSelectStatement(Constants.NO_SELECT_TIMEOUT);  // No timeout on metadata queries
             long start = System.currentTimeMillis();
             rs = stmt.executeQuery(query);
             logger.trace("[{} ms] {}", (System.currentTimeMillis() - start), query);

diff --git a/src/main/java/com/feedzai/commons/sql/abstraction/engine/DatabaseEngine.java b/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.
      *

diff --git a/src/main/java/com/feedzai/commons/sql/abstraction/engine/DatabaseEngineTimeoutException.java b/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);
+    }
+}
diff --git a/src/main/java/com/feedzai/commons/sql/abstraction/engine/configuration/PdbProperties.java b/src/main/java/com/feedzai/commons/sql/abstraction/engine/configuration/PdbProperties.java
@@ -33,6 +33,7 @@
 import static com.feedzai.commons.sql.abstraction.util.Constants.DEFAULT_LOGIN_TIMEOUT;
 import static com.feedzai.commons.sql.abstraction.util.Constants.DEFAULT_MAXIMUM_TIME_BATCH_SHUTDOWN;
 import static com.feedzai.commons.sql.abstraction.util.Constants.DEFAULT_MAX_IDENTIFIER_SIZE;
+import static com.feedzai.commons.sql.abstraction.util.Constants.DEFAULT_SELECT_QUERY_TIMEOUT;
 import static com.feedzai.commons.sql.abstraction.util.Constants.DEFAULT_RECONNECT_ON_LOST;
 import static com.feedzai.commons.sql.abstraction.util.Constants.DEFAULT_RETRY_INTERVAL;
 import static com.feedzai.commons.sql.abstraction.util.Constants.DEFAULT_SCHEMA_POLICY;
@@ -166,6 +167,12 @@ public class PdbProperties extends Properties implements com.feedzai.commons.sql
      */
     public static final String SOCKET_TIMEOUT = "pdb.socket_timeout";
 
+    /**
+     * Property that indicates the maximum time a select query is allowed to run, in seconds. No limit if zero.
+     * See {@link java.sql.Statement#setQueryTimeout(int)}.
+     */
+    public static final String SELECT_QUERY_TIMEOUT = "pdb.query_select_timeout";
+
     /**
      * Creates a new instance of an empty {@link PdbProperties}.
      */
@@ -199,6 +206,7 @@ public PdbProperties(final boolean useDefaults) {
             setProperty(DISABLE_LOB_CACHING, DEFAULT_DISABLE_LOB_CACHING);
             setProperty(LOGIN_TIMEOUT, DEFAULT_LOGIN_TIMEOUT);
             setProperty(SOCKET_TIMEOUT, DEFAULT_SOCKET_TIMEOUT);
+            setProperty(SELECT_QUERY_TIMEOUT, DEFAULT_SELECT_QUERY_TIMEOUT);
         }
     }
 
@@ -495,6 +503,15 @@ public int getSocketTimeout() {
         return Integer.parseInt(getProperty(SOCKET_TIMEOUT));
     }
 
+    /**
+     * Gets the query select timeout (in seconds).
+     *
+     * @return The query select timeout.
+     */
+    public int getSelectQueryTimeout() {
+        return Integer.parseInt(getProperty(SELECT_QUERY_TIMEOUT));
+    }
+
     /**
      * Checks if the driver property is set.
      *

diff --git a/src/main/java/com/feedzai/commons/sql/abstraction/util/Constants.java b/src/main/java/com/feedzai/commons/sql/abstraction/util/Constants.java
@@ -96,4 +96,15 @@ public final class Constants {
      * By default, there is no timeout at the socket level, so pdb will wait indefinitely for the database to respond the queries.
      */
     public static final int DEFAULT_SOCKET_TIMEOUT = 0;
+
+    /**
+     * The value that represents absence of a timeout in query timeouts.
+     */
+    public static final int NO_SELECT_TIMEOUT = 0;
+
+    /**
+     * The default select query timeout (in seconds).
+     * By default, there is no query timeout, so pdb will wait indefinitely for the database to respond to queries.
+     */
+    public static final int DEFAULT_SELECT_QUERY_TIMEOUT = NO_SELECT_TIMEOUT;
 }