From ec6891a08e4fc2f775d9ad4387f6db1da65ccec4 Mon Sep 17 00:00:00 2001
From: Brett Wooldridge <brett.wooldridge@gmail.com>
Date: Sun, 17 Nov 2024 02:45:38 +0900
Subject: [PATCH] Expand/improve SQLExceptionOverride handling flexibility

---
 .../zaxxer/hikari/SQLExceptionOverride.java   | 25 +++++++++++--------
 .../zaxxer/hikari/pool/ProxyConnection.java   |  9 ++++---
 2 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/src/main/java/com/zaxxer/hikari/SQLExceptionOverride.java b/src/main/java/com/zaxxer/hikari/SQLExceptionOverride.java
index 17168d8e..84b8dfe6 100644
--- a/src/main/java/com/zaxxer/hikari/SQLExceptionOverride.java
+++ b/src/main/java/com/zaxxer/hikari/SQLExceptionOverride.java
@@ -4,24 +4,29 @@ 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.
+ * of HikariCP. When a SQLException is thrown from JDBC execution methods, the
+ * SQLState and error code will be checked to determine if the connection should
+ * be evicted from the pool.
+ * <p>
+ * By supplying an implementation of this interface, users can override the default
+ * handling of SQLExceptions. The {@link #adjudicate(SQLException)} method will be called
+ * with the SQLException that was thrown. If the method returns {@link Override#CONTINUE_EVICT}
+ * the customary built-in handling will occur. If the method returns {@link Override#DO_NOT_EVICT}
+ * the eviction will be elided. If the method returns {@link Override#MUST_EVICT} the eviction will
+ * be evicted regardless of the SQLState or error code.
  */
 public interface SQLExceptionOverride {
    enum Override {
       CONTINUE_EVICT,
-      DO_NOT_EVICT
+      DO_NOT_EVICT,
+      MUST_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.
+    * This method is called when a SQLException is thrown from a JDBC method.
     *
-    * @param sqlException the #SQLException to adjudicate
-    * @return either one of {@link Override#CONTINUE_EVICT} or {@link Override#DO_NOT_EVICT}
+    * @param sqlException the SQLException that was thrown
+    * @return an {@link Override} value indicating how eviction should proceed
     */
    default Override adjudicate(final SQLException sqlException)
    {
diff --git a/src/main/java/com/zaxxer/hikari/pool/ProxyConnection.java b/src/main/java/com/zaxxer/hikari/pool/ProxyConnection.java
index b6dd3f54..02526d0c 100644
--- a/src/main/java/com/zaxxer/hikari/pool/ProxyConnection.java
+++ b/src/main/java/com/zaxxer/hikari/pool/ProxyConnection.java
@@ -27,8 +27,7 @@ import java.util.HashSet;
 import java.util.Set;
 import java.util.concurrent.Executor;
 
-import static com.zaxxer.hikari.SQLExceptionOverride.Override.CONTINUE_EVICT;
-import static com.zaxxer.hikari.SQLExceptionOverride.Override.DO_NOT_EVICT;
+import static com.zaxxer.hikari.SQLExceptionOverride.Override.*;
 
 /**
  * This is the proxy class for {@link Connection}.
@@ -156,12 +155,14 @@ public abstract class ProxyConnection implements Connection
       final var exceptionOverride = poolEntry.getPoolBase().exceptionOverride;
       for (int depth = 0; delegate != ClosedConnection.CLOSED_CONNECTION && nse != null && depth < 10; depth++) {
          final var sqlState = nse.getSQLState();
-         if (exceptionOverride != null && exceptionOverride.adjudicate(nse) == DO_NOT_EVICT) {
+         final var shouldEvict = exceptionOverride != null ? exceptionOverride.adjudicate(nse) : CONTINUE_EVICT;
+         if (shouldEvict == DO_NOT_EVICT) {
             break;
          }
          else if (sqlState != null && sqlState.startsWith("08")
              || ERROR_STATES.contains(sqlState)
-             || ERROR_CODES.contains(nse.getErrorCode())) {
+             || ERROR_CODES.contains(nse.getErrorCode())
+             || shouldEvict == MUST_EVICT) {
 
             // broken connection
             evict = true;