Fixes #1489 allow specification of a user supplied exception eviction override class

src/main/java/com/zaxxer/hikari/HikariConfig.java

@@ -80,6 +80,7 @@ public class HikariConfig implements HikariConfigMXBean
    private String dataSourceClassName;
    private String dataSourceJndiName;
    private String driverClassName;
+   private String exceptionOverrideClassName;
    private String jdbcUrl;
    private String poolName;
    private String schema;
@@ -812,7 +813,8 @@ public class HikariConfig implements HikariConfigMXBean
     *
     * @return the default schema name
     */
-   public String getSchema() {
+   public String getSchema()
+   {
       return schema;
    }
 
@@ -827,6 +829,50 @@ public class HikariConfig implements HikariConfigMXBean
       this.schema = schema;
    }
 
+   /**
+    * Get the user supplied SQLExceptionOverride class name.
+    *
+    * @return the user supplied SQLExceptionOverride class name
+    * @see SQLExceptionOverride
+    */
+   public String getExceptionOverrideClassName()
+   {
+      return this.exceptionOverrideClassName;
+   }
+
+   /**
+    * Set the user supplied SQLExceptionOverride class name.
+    *
+    * @param exceptionOverrideClassName the user supplied SQLExceptionOverride class name
+    * @see SQLExceptionOverride
+    */
+   public void setExceptionOverrideClassName(String exceptionOverrideClassName)
+   {
+      checkIfSealed();
+
+      Class<?> overrideClass = attemptFromContextLoader(exceptionOverrideClassName);
+      try {
+         if (overrideClass == null) {
+            overrideClass = this.getClass().getClassLoader().loadClass(exceptionOverrideClassName);
+            LOGGER.debug("SQLExceptionOverride class {} found in the HikariConfig class classloader {}", exceptionOverrideClassName, this.getClass().getClassLoader());
+         }
+      } catch (ClassNotFoundException e) {
+         LOGGER.error("Failed to load SQLExceptionOverride class {} from HikariConfig class classloader {}", exceptionOverrideClassName, this.getClass().getClassLoader());
+      }
+
+      if (overrideClass == null) {
+         throw new RuntimeException("Failed to load SQLExceptionOverride class " + exceptionOverrideClassName + " in either of HikariConfig class loader or Thread context classloader");
+      }
+
+      try {
+         overrideClass.getConstructor().newInstance();
+         this.exceptionOverrideClassName = exceptionOverrideClassName;
+      }
+      catch (Exception e) {
+         throw new RuntimeException("Failed to instantiate class " + exceptionOverrideClassName, e);
+      }
+   }
+
    /**
     * Set the default transaction isolation level.  The specified value is the
     * constant name from the <code>Connection</code> class, eg.

src/main/java/com/zaxxer/hikari/SQLExceptionOverride.java

@@ -0,0 +1,30 @@
+package com.zaxxer.hikari;
+
+import java.sql.SQLException;
+
+/**
+ * Users can implement this interface to override the default SQLException handling
+ * of HikariCP.  By the time an instance of this interface is invoked HikariCP has
+ * already made a determination to evict the Connection from the pool.
+ *
+ * If the {@link #adjudicate(SQLException)} method returns {@link Override#CONTINUE_EVICT} the eviction will occur, but if the
+ * method returns {@link Override#DO_NOT_EVICT} the eviction will be elided.
+ */
+public interface SQLExceptionOverride {
+   enum Override {
+      CONTINUE_EVICT,
+      DO_NOT_EVICT
+   }
+
+   /**
+    * If this method returns {@link Override#CONTINUE_EVICT} then Connection eviction will occur, but if it
+    * returns {@link Override#DO_NOT_EVICT} the eviction will be elided.
+    *
+    * @param sqlException the #SQLException to adjudicate
+    * @return either one of {@link Override#CONTINUE_EVICT} or {@link Override#DO_NOT_EVICT}
+    */
+   default Override adjudicate(final SQLException sqlException)
+   {
+      return Override.CONTINUE_EVICT;
+   }
+}

src/main/java/com/zaxxer/hikari/pool/PoolBase.java

@@ -17,6 +17,7 @@
 package com.zaxxer.hikari.pool;
 
 import com.zaxxer.hikari.HikariConfig;
+import com.zaxxer.hikari.SQLExceptionOverride;
 import com.zaxxer.hikari.metrics.IMetricsTracker;
 import com.zaxxer.hikari.pool.HikariPool.PoolInitializationException;
 import com.zaxxer.hikari.util.DriverDataSource;
@@ -64,6 +65,8 @@ abstract class PoolBase
    long connectionTimeout;
    long validationTimeout;
 
+   SQLExceptionOverride exceptionOverride;
+
    private static final String[] RESET_STATES = {"readOnly", "autoCommit", "isolation", "catalog", "netTimeout", "schema"};
    private static final int UNINITIALIZED = -1;
    private static final int TRUE = 1;
@@ -95,6 +98,7 @@ abstract class PoolBase
       this.schema = config.getSchema();
       this.isReadOnly = config.isReadOnly();
       this.isAutoCommit = config.isAutoCommit();
+      this.exceptionOverride = UtilityElf.createInstance(config.getExceptionOverrideClassName(), SQLExceptionOverride.class);
       this.transactionIsolation = UtilityElf.getTransactionIsolation(config.getTransactionIsolation());
 
       this.isQueryTimeoutSupported = UNINITIALIZED;

src/main/java/com/zaxxer/hikari/pool/PoolEntry.java

@@ -128,6 +128,11 @@ final class PoolEntry implements IConcurrentBagEntry
       return elapsedMillis(lastBorrowed);
    }
 
+   PoolBase getPoolBase()
+   {
+      return hikariPool;
+   }
+
    /** {@inheritDoc} */
    @Override
    public String toString()

src/main/java/com/zaxxer/hikari/pool/ProxyConnection.java

@@ -16,6 +16,7 @@
 
 package com.zaxxer.hikari.pool;
 
+import com.zaxxer.hikari.SQLExceptionOverride;
 import com.zaxxer.hikari.util.FastList;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -27,6 +28,7 @@ import java.util.HashSet;
 import java.util.Set;
 import java.util.concurrent.Executor;
 
+import static com.zaxxer.hikari.SQLExceptionOverride.Override.DO_NOT_EVICT;
 import static com.zaxxer.hikari.util.ClockSource.currentTime;
 
 /**
@@ -83,7 +85,13 @@ public abstract class ProxyConnection implements Connection
       ERROR_CODES.add(2399);
    }
 
-   protected ProxyConnection(final PoolEntry poolEntry, final Connection connection, final FastList<Statement> openStatements, final ProxyLeakTask leakTask, final long now, final boolean isReadOnly, final boolean isAutoCommit) {
+   protected ProxyConnection(final PoolEntry poolEntry,
+                             final Connection connection,
+                             final FastList<Statement> openStatements,
+                             final ProxyLeakTask leakTask,
+                             final long now,
+                             final boolean isReadOnly,
+                             final boolean isAutoCommit) {
       this.poolEntry = poolEntry;
       this.delegate = connection;
       this.openStatements = openStatements;
@@ -143,9 +151,12 @@ public abstract class ProxyConnection implements Connection
       return poolEntry;
    }
 
+   @SuppressWarnings("ConstantConditions")
    final SQLException checkException(SQLException sqle)
    {
+      boolean evict = false;
       SQLException nse = sqle;
+      final SQLExceptionOverride exceptionOverride = poolEntry.getPoolBase().exceptionOverride;
       for (int depth = 0; delegate != ClosedConnection.CLOSED_CONNECTION && nse != null && depth < 10; depth++) {
          final String sqlState = nse.getSQLState();
          if (sqlState != null && sqlState.startsWith("08")
@@ -153,18 +164,28 @@ public abstract class ProxyConnection implements Connection
              || ERROR_STATES.contains(sqlState)
              || ERROR_CODES.contains(nse.getErrorCode())) {
 
+            if (exceptionOverride != null && exceptionOverride.adjudicate(nse) == DO_NOT_EVICT) {
+               break;
+            }
+
             // broken connection
-            LOGGER.warn("{} - Connection {} marked as broken because of SQLSTATE({}), ErrorCode({})",
+            evict = true;
-                        poolEntry.getPoolName(), delegate, sqlState, nse.getErrorCode(), nse);
+            break;
-            leakTask.cancel();
-            poolEntry.evict("(connection is broken)");
-            delegate = ClosedConnection.CLOSED_CONNECTION;
          }
          else {
             nse = nse.getNextException();
          }
       }
 
+      if (evict) {
+         SQLException exception = (nse != null) ? nse : sqle;
+         LOGGER.warn("{} - Connection {} marked as broken because of SQLSTATE({}), ErrorCode({})",
+            poolEntry.getPoolName(), delegate, exception.getSQLState(), exception.getErrorCode(), exception);
+         leakTask.cancel();
+         poolEntry.evict("(connection is broken)");
+         delegate = ClosedConnection.CLOSED_CONNECTION;
+      }
+
       return sqle;
    }
 

src/test/java/com/zaxxer/hikari/pool/TestConnections.java

@@ -19,6 +19,7 @@ package com.zaxxer.hikari.pool;
 import com.zaxxer.hikari.HikariConfig;
 import com.zaxxer.hikari.HikariDataSource;
 import com.zaxxer.hikari.HikariPoolMXBean;
+import com.zaxxer.hikari.SQLExceptionOverride;
 import com.zaxxer.hikari.mocks.StubConnection;
 import com.zaxxer.hikari.mocks.StubDataSource;
 import com.zaxxer.hikari.mocks.StubStatement;
@@ -40,6 +41,7 @@ import static org.junit.Assert.*;
 /**
  * @author Brett Wooldridge
  */
+@SuppressWarnings({"SqlDialectInspection", "SqlNoDataSourceInspection"})
 public class TestConnections
 {
    @Before
@@ -258,6 +260,81 @@ public class TestConnections
       }
    }
 
+   @Test
+   public void testEviction2() throws SQLException
+   {
+      HikariConfig config = newHikariConfig();
+      config.setMaximumPoolSize(5);
+      config.setConnectionTimeout(2500);
+      config.setConnectionTestQuery("VALUES 1");
+      config.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");
+      config.setExceptionOverrideClassName(OverrideHandler.class.getName());
+
+      try (HikariDataSource ds = new HikariDataSource(config)) {
+         HikariPool pool = getPool(ds);
+
+         while (pool.getTotalConnections() < 5) {
+            quietlySleep(100L);
+         }
+
+         try (Connection connection = ds.getConnection()) {
+            assertNotNull(connection);
+
+            PreparedStatement statement = connection.prepareStatement("SELECT some, thing FROM somewhere WHERE something=?");
+            assertNotNull(statement);
+
+            ResultSet resultSet = statement.executeQuery();
+            assertNotNull(resultSet);
+
+            try {
+               statement.getMaxFieldSize();
+            } catch (Exception e) {
+               assertSame(SQLException.class, e.getClass());
+            }
+         }
+
+         assertEquals("Total connections not as expected", 5, pool.getTotalConnections());
+         assertEquals("Idle connections not as expected", 5, pool.getIdleConnections());
+      }
+   }
+
+   @Test
+   public void testEviction3() throws SQLException
+   {
+      HikariConfig config = newHikariConfig();
+      config.setMaximumPoolSize(5);
+      config.setConnectionTimeout(2500);
+      config.setConnectionTestQuery("VALUES 1");
+      config.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");
+
+      try (HikariDataSource ds = new HikariDataSource(config)) {
+         HikariPool pool = getPool(ds);
+
+         while (pool.getTotalConnections() < 5) {
+            quietlySleep(100L);
+         }
+
+         try (Connection connection = ds.getConnection()) {
+            assertNotNull(connection);
+
+            PreparedStatement statement = connection.prepareStatement("SELECT some, thing FROM somewhere WHERE something=?");
+            assertNotNull(statement);
+
+            ResultSet resultSet = statement.executeQuery();
+            assertNotNull(resultSet);
+
+            try {
+               statement.getMaxFieldSize();
+            } catch (Exception e) {
+               assertSame(SQLException.class, e.getClass());
+            }
+         }
+
+         assertEquals("Total connections not as expected", 4, pool.getTotalConnections());
+         assertEquals("Idle connections not as expected", 4, pool.getIdleConnections());
+      }
+   }
+
    @Test
    public void testEvictAllRefill() throws Exception {
       HikariConfig config = newHikariConfig();
@@ -735,12 +812,13 @@ public class TestConnections
       }
    }
 
-   class StubDataSourceWithErrorSwitch extends StubDataSource {
+   static class StubDataSourceWithErrorSwitch extends StubDataSource
+   {
       private boolean errorOnConnection = false;
 
       /** {@inheritDoc} */
       @Override
-      public Connection getConnection() throws SQLException {
+      public Connection getConnection() {
          if (!errorOnConnection) {
             return new StubConnection();
          }
@@ -753,4 +831,11 @@ public class TestConnections
       }
    }
 
+   public static class OverrideHandler implements SQLExceptionOverride
+   {
+      @java.lang.Override
+      public Override adjudicate(SQLException sqlException) {
+         return (sqlException.getSQLState().equals("08999")) ? Override.DO_NOT_EVICT : Override.CONTINUE_EVICT;
+      }
+   }
 }