Return the connection when withConnection throws

vertx-pg-client/src/test/java/io/vertx/tests/pgclient/PgPoolTestBase.java

@@ -30,6 +30,7 @@
 import org.junit.Before;
 import org.junit.Test;
 
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 
@@ -256,4 +257,24 @@ public void testWithConnection(TestContext ctx) {
       }
     }
   }
+
+  @Test
+  public void testWithConnectionFunctionThrows(TestContext ctx) {
+    Pool pool = createPool(options, new PoolOptions()
+      .setMaxSize(1)
+      .setConnectionTimeout(2)
+      .setConnectionTimeoutUnit(TimeUnit.SECONDS));
+    RuntimeException failure = new RuntimeException("The function fails");
+    pool
+      .withConnection(conn -> {
+        throw failure;
+      })
+      .onComplete(ctx.asyncAssertFailure(err -> {
+        ctx.assertEquals(failure, err);
+        // The only connection of the pool must have been returned
+        pool
+          .withConnection(conn -> conn.query("SELECT 1").execute())
+          .onComplete(ctx.asyncAssertSuccess());
+      }));
+  }
 }

vertx-sql-client/src/main/java/io/vertx/sqlclient/Pool.java

@@ -185,7 +185,15 @@ static Pool pool(Vertx vertx, SqlConnectOptions database, PoolOptions options) {
    * @return a future notified with the result
    */
   default <T> Future<@Nullable T> withConnection(Function<SqlConnection, Future<@Nullable T>> function) {
-    return getConnection().flatMap(conn -> function.apply(conn).onComplete(ar -> conn.close()));
+    return getConnection().flatMap(conn -> {
+      Future<@Nullable T> result;
+      try {
+        result = function.apply(conn);
+      } catch (Throwable t) {
+        result = Future.failedFuture(t);
+      }
+      return result.onComplete(ar -> conn.close());
+    });
   }
 
   /**