Simplify handling of propagated connections, place connection in context

Author: technical-debt-collector

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

@@ -98,12 +98,6 @@ static Pool pool(Vertx vertx, SqlConnectOptions database, PoolOptions options) {
    */
   Future<SqlConnection> getConnection();
 
-  Future<SqlConnection> getPropagatableConnection();
-
-  boolean propagatableConnectionIsActive();
-
-  Future<Void> setPropagatableConnection(SqlConnection propagatableConnection);
-
   /**
    * {@inheritDoc}
    *
@@ -170,38 +164,11 @@ default <T> void withTransaction(Function<SqlConnection, Future<@Nullable T>> fu
         .onComplete(ar -> conn.close()));
   }
 
-
-  default <T> Future<@Nullable T> withPropagatedTransaction(Function<SqlConnection, Future<@Nullable T>> function) {
-    if (propagatableConnectionIsActive()) {
-      return getPropagatableConnection()
-        .flatMap(conn -> function.apply(conn)
-          .onFailure(err -> {
-            if (!(err instanceof TransactionRollbackException)) {
-              conn.getTransaction().rollback();
-            }
-          }));
-    } else {
-      return getPropagatableConnection()
-        .flatMap(conn -> conn
-          .begin()
-          .flatMap(tx -> function
-            .apply(conn)
-            .compose(
-              res -> tx
-                .commit()
-                .flatMap(v -> Future.succeededFuture(res)),
-              err -> {
-                if (err instanceof TransactionRollbackException) {
-                  return Future.failedFuture(err);
-                } else {
-                  return tx
-                    .rollback()
-                    .compose(v -> Future.failedFuture(err), failure -> Future.failedFuture(err));
-                }
-              }))
-          .onComplete(ar -> conn.close(v -> setPropagatableConnection(null))));
-    }
-  }
+  /**
+   * Like {@link #withTransaction(Function, Handler)} but keeps the connection accessible via the context
+   * for the duration of the given {@code function}
+   */
+  <T> Future<@Nullable T> withPropagatedTransaction(Function<SqlConnection, Future<@Nullable T>> function);
 
   /**
    * Get a connection from the pool and execute the given {@code function}.

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

@@ -15,6 +15,7 @@
  */
 package io.vertx.sqlclient.impl;
 
+import io.vertx.codegen.annotations.Nullable;
 import io.vertx.core.AsyncResult;
 import io.vertx.core.Context;
 import io.vertx.core.Future;
@@ -67,21 +68,6 @@ public Future<SqlConnection> getConnection() {
     return delegate.getConnection();
   }
 
-  @Override
-  public Future<SqlConnection> getPropagatableConnection() {
-    return delegate.getPropagatableConnection();
-  }
-
-  @Override
-  public boolean propagatableConnectionIsActive() {
-    return delegate.propagatableConnectionIsActive();
-  }
-
-  @Override
-  public Future<Void> setPropagatableConnection(SqlConnection propagatableConnection) {
-    return delegate.setPropagatableConnection(propagatableConnection);
-  }
-
   @Override
   public Query<RowSet<Row>> query(String sql) {
     return delegate.query(sql);
@@ -92,6 +78,11 @@ public PreparedQuery<RowSet<Row>> preparedQuery(String sql) {
     return delegate.preparedQuery(sql);
   }
 
+  @Override
+  public <T> Future<@Nullable T> withPropagatedTransaction(Function<SqlConnection, Future<@Nullable T>> function) {
+    return delegate.withPropagatedTransaction(function);
+  }
+
   @Override
   public P connectHandler(Handler<SqlConnection> handler) {
     delegate.connectHandler(handler);

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

@@ -17,6 +17,7 @@
 
 package io.vertx.sqlclient.impl;
 
+import io.vertx.codegen.annotations.Nullable;
 import io.vertx.core.*;
 import io.vertx.core.impl.CloseFuture;
 import io.vertx.core.impl.ContextInternal;
@@ -26,12 +27,12 @@
 import io.vertx.sqlclient.Pool;
 import io.vertx.sqlclient.PoolOptions;
 import io.vertx.sqlclient.SqlConnection;
+import io.vertx.sqlclient.TransactionRollbackException;
 import io.vertx.sqlclient.impl.command.CommandBase;
 import io.vertx.sqlclient.impl.pool.SqlConnectionPool;
 import io.vertx.sqlclient.impl.tracing.QueryTracer;
 import io.vertx.sqlclient.spi.Driver;
 
-import java.util.Objects;
 import java.util.function.Function;
 
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
@@ -53,7 +54,7 @@ public class PoolImpl extends SqlClientBase implements Pool, Closeable {
   private long timerID;
   private volatile Function<Context, Future<SqlConnection>> connectionProvider;
 
-  private volatile SqlConnection propagatableConnection;
+  private static final String PROPAGATABLE_CONNECTION = "propagatable_connection";
 
   public PoolImpl(VertxInternal vertx,
                   Driver driver,
@@ -157,22 +158,40 @@ public Future<SqlConnection> getConnection() {
     });
   }
 
-  @Override
-  public Future<SqlConnection> getPropagatableConnection() {
-    if (propagatableConnection == null) {
-      return getConnection().onComplete(c -> setPropagatableConnection(c.result()));
+  public <T> Future<@Nullable T> withPropagatedTransaction(Function<SqlConnection, Future<@Nullable T>> function) {
+    SqlConnection sqlConnection = context().get(PROPAGATABLE_CONNECTION);
+    if (sqlConnection == null) {
+      return initializePropagatedConnectionAndTransaction(function);
     }
-    return Future.succeededFuture(propagatableConnection);
+    return Future.succeededFuture(sqlConnection)
+      .flatMap(conn -> function.apply(conn)
+        .onFailure(err -> {
+          if (!(err instanceof TransactionRollbackException)) {
+            conn.getTransaction().rollback();
+          }
+        }));
   }
 
-  @Override
-  public boolean propagatableConnectionIsActive() {
-    return propagatableConnection != null;
-  }
-
-  @Override
-  public Future<Void> setPropagatableConnection(SqlConnection propagatableConnection) {
-    return Future.future(handler -> this.propagatableConnection = propagatableConnection);
+  private <T> Future<@Nullable T> initializePropagatedConnectionAndTransaction(Function<SqlConnection, Future<@Nullable T>> function) {
+    return getConnection().onComplete(handler -> context().put(PROPAGATABLE_CONNECTION, handler.result()))
+      .flatMap(conn -> conn
+        .begin()
+        .flatMap(tx -> function
+          .apply(conn)
+          .compose(
+            res -> tx
+              .commit()
+              .flatMap(v -> Future.succeededFuture(res)),
+            err -> {
+              if (err instanceof TransactionRollbackException) {
+                return Future.failedFuture(err);
+              } else {
+                return tx
+                  .rollback()
+                  .compose(v -> Future.failedFuture(err), failure -> Future.failedFuture(err));
+              }
+            }))
+        .onComplete(ar -> conn.close(v -> context().remove(PROPAGATABLE_CONNECTION))));
   }
 
   @Override

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

@@ -321,7 +321,6 @@ public void testWithPropagatedConnectionTransactionCommit(TestContext ctx) {
       .query("SELECT id, val FROM mutable")
       .execute(ctx.asyncAssertSuccess(rows -> {
         ctx.assertEquals(3, rows.size());
-        ctx.assertFalse(pool.propagatableConnectionIsActive());
         async.complete();
       }))));
   }