Associate futures under propagated connection handling with the current context

Author: Ola Sæter Heitmann

vertx-sql-client/src/main/java/io/vertx/sqlclient/impl/PoolImpl.java

@@ -159,12 +159,12 @@ public Future<SqlConnection> getConnection() {
   }
 
   public <T> Future<@Nullable T> withPropagatedTransaction(Function<SqlConnection, Future<@Nullable T>> function) {
-    Context context = Vertx.currentContext();
+    ContextInternal context = (ContextInternal) Vertx.currentContext();
     SqlConnection sqlConnection = context.getLocal(PROPAGATABLE_CONNECTION);
     if (sqlConnection == null) {
       return initializePropagatedConnectionAndTransaction(function);
     }
-    return Future.succeededFuture(sqlConnection)
+    return context.succeededFuture(sqlConnection)
       .flatMap(conn -> function.apply(conn)
         .onFailure(err -> {
           if (!(err instanceof TransactionRollbackException)) {
@@ -174,7 +174,7 @@ public Future<SqlConnection> getConnection() {
   }
 
   private <T> Future<@Nullable T> initializePropagatedConnectionAndTransaction(Function<SqlConnection, Future<@Nullable T>> function) {
-    Context context = Vertx.currentContext();
+    ContextInternal context = (ContextInternal) Vertx.currentContext();
     return getConnection().onComplete(handler -> context.putLocal(PROPAGATABLE_CONNECTION, handler.result()))
       .flatMap(conn -> conn
         .begin()
@@ -183,14 +183,14 @@ public Future<SqlConnection> getConnection() {
           .compose(
             res -> tx
               .commit()
-              .flatMap(v -> Future.succeededFuture(res)),
+              .flatMap(v -> context.succeededFuture(res)),
             err -> {
               if (err instanceof TransactionRollbackException) {
-                return Future.failedFuture(err);
+                return context.failedFuture(err);
               } else {
                 return tx
                   .rollback()
-                  .compose(v -> Future.failedFuture(err), failure -> Future.failedFuture(err));
+                  .compose(v -> context.failedFuture(err), failure -> context.failedFuture(err));
               }
             }))
         .onComplete(ar -> conn.close(v -> context.removeLocal(PROPAGATABLE_CONNECTION))));