Create propagatable connection, add Transaction accessor

Author: technical-debt-collector

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

@@ -98,6 +98,12 @@ static Pool pool(Vertx vertx, SqlConnectOptions database, PoolOptions options) {
    */
   Future<SqlConnection> getConnection();
 
+  Future<SqlConnection> getPropagatableConnection();
+
+  boolean propagatableConnectionIsActive();
+
+  Future<Void> setPropagatableConnection(SqlConnection propagatableConnection);
+
   /**
    * {@inheritDoc}
    *
@@ -164,6 +170,39 @@ 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))));
+    }
+  }
+
   /**
    * Get a connection from the pool and execute the given {@code function}.
    *

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

@@ -92,6 +92,8 @@ public interface SqlConnection extends SqlClient {
    */
   Future<Transaction> begin();
 
+  Transaction getTransaction();
+
   /**
    * @return whether the connection uses SSL
    */

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

@@ -67,6 +67,21 @@ 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);

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

@@ -53,6 +53,8 @@ public class PoolImpl extends SqlClientBase implements Pool, Closeable {
   private long timerID;
   private volatile Function<Context, Future<SqlConnection>> connectionProvider;
 
+  private volatile SqlConnection propagatableConnection;
+
   public PoolImpl(VertxInternal vertx,
                   Driver driver,
                   QueryTracer tracer,
@@ -155,6 +157,24 @@ public Future<SqlConnection> getConnection() {
     });
   }
 
+  @Override
+  public Future<SqlConnection> getPropagatableConnection() {
+    if (propagatableConnection == null) {
+      return getConnection().onComplete(c -> setPropagatableConnection(c.result()));
+    }
+    return Future.succeededFuture(propagatableConnection);
+  }
+
+  @Override
+  public boolean propagatableConnectionIsActive() {
+    return propagatableConnection != null;
+  }
+
+  @Override
+  public Future<Void> setPropagatableConnection(SqlConnection propagatableConnection) {
+    return Future.future(handler -> this.propagatableConnection = propagatableConnection);
+  }
+
   @Override
   public <R> Future<R> schedule(ContextInternal context, CommandBase<R> cmd) {
     Object metric;

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

@@ -165,6 +165,11 @@ public Future<Transaction> begin() {
     return tx.begin();
   }
 
+  @Override
+  public Transaction getTransaction() {
+    return tx;
+  }
+
   @Override
   boolean autoCommit() {
     return tx == null;

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

@@ -308,4 +308,36 @@ public void testWithTransactionImplicitRollback(TestContext ctx) {
         }));
     }));
   }
+
+  @Test
+  public void testWithPropagatedConnectionTransactionCommit(TestContext ctx) {
+    Async async = ctx.async();
+    Pool pool = createPool();
+    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 ->
+        c.query("INSERT INTO mutable (id, val) VALUES (3, 'hello-3')").execute().mapEmpty())
+    ).onComplete(ctx.asyncAssertSuccess(v -> pool
+      .query("SELECT id, val FROM mutable")
+      .execute(ctx.asyncAssertSuccess(rows -> {
+        ctx.assertEquals(3, rows.size());
+        ctx.assertFalse(pool.propagatableConnectionIsActive());
+        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();
+      }))));
+  }
 }