Use current context, store connection locally

Author: technical-debt-collector

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

@@ -159,7 +159,8 @@ public Future<SqlConnection> getConnection() {
   }
 
   public <T> Future<@Nullable T> withPropagatedTransaction(Function<SqlConnection, Future<@Nullable T>> function) {
-    SqlConnection sqlConnection = context().get(PROPAGATABLE_CONNECTION);
+    Context context = Vertx.currentContext();
+    SqlConnection sqlConnection = context.getLocal(PROPAGATABLE_CONNECTION);
     if (sqlConnection == null) {
       return initializePropagatedConnectionAndTransaction(function);
     }
@@ -173,7 +174,8 @@ public Future<SqlConnection> getConnection() {
   }
 
   private <T> Future<@Nullable T> initializePropagatedConnectionAndTransaction(Function<SqlConnection, Future<@Nullable T>> function) {
-    return getConnection().onComplete(handler -> context().put(PROPAGATABLE_CONNECTION, handler.result()))
+    Context context = Vertx.currentContext();
+    return getConnection().onComplete(handler -> context.putLocal(PROPAGATABLE_CONNECTION, handler.result()))
       .flatMap(conn -> conn
         .begin()
         .flatMap(tx -> function
@@ -191,7 +193,7 @@ public Future<SqlConnection> getConnection() {
                   .compose(v -> Future.failedFuture(err), failure -> Future.failedFuture(err));
               }
             }))
-        .onComplete(ar -> conn.close(v -> context().remove(PROPAGATABLE_CONNECTION))));
+        .onComplete(ar -> conn.close(v -> context.removeLocal(PROPAGATABLE_CONNECTION))));
   }
 
   @Override

vertx-sql-client/src/test/java/io/vertx/sqlclient/tck/TransactionTestBase.java

@@ -313,6 +313,7 @@ public void testWithTransactionImplicitRollback(TestContext ctx) {
   public void testWithPropagatedConnectionTransactionCommit(TestContext ctx) {
     Async async = ctx.async();
     Pool pool = createPool();
+    vertx.runOnContext(handler -> {
     pool.withPropagatedTransaction(c ->
       pool.withPropagatedTransaction(conn -> conn.query("INSERT INTO mutable (id, val) VALUES (1, 'hello-1')").execute().mapEmpty()).flatMap(v ->
         pool.withPropagatedTransaction(conn -> conn.query("INSERT INTO mutable (id, val) VALUES (2, 'hello-2')").execute().mapEmpty())).flatMap(v2 ->
@@ -321,22 +322,27 @@ public void testWithPropagatedConnectionTransactionCommit(TestContext ctx) {
       .query("SELECT id, val FROM mutable")
       .execute(ctx.asyncAssertSuccess(rows -> {
         ctx.assertEquals(3, rows.size());
+        ctx.assertNull(Vertx.currentContext().getLocal("propagatable_connection"));
         async.complete();
       }))));
+    });
   }
 
   @Test
   public void testWithPropagatedConnectionTransactionRollback(TestContext ctx) {
     Async async = ctx.async();
     Pool pool = createPool();
     Throwable failure = new Throwable();
-    pool.withPropagatedTransaction(c ->
-      pool.withPropagatedTransaction(conn -> conn.query("INSERT INTO mutable (id, val) VALUES (1, 'hello-1')").execute().mapEmpty().flatMap(v -> Future.failedFuture(failure)))
-    ).onComplete(ctx.asyncAssertFailure(v -> pool
-      .query("SELECT id, val FROM mutable")
-      .execute(ctx.asyncAssertSuccess(rows -> {
-        ctx.assertEquals(0, rows.size());
-        async.complete();
-      }))));
+    vertx.runOnContext(handler -> {
+      pool.withPropagatedTransaction(c ->
+        pool.withPropagatedTransaction(conn -> conn.query("INSERT INTO mutable (id, val) VALUES (1, 'hello-1')").execute().mapEmpty().flatMap(v -> Future.failedFuture(failure)))
+      ).onComplete(ctx.asyncAssertFailure(v -> pool
+        .query("SELECT id, val FROM mutable")
+        .execute(ctx.asyncAssertSuccess(rows -> {
+          ctx.assertEquals(0, rows.size());
+          ctx.assertNull(Vertx.currentContext().getLocal("propagatable_connection"));
+          async.complete();
+        }))));
+    });
   }
 }