Add support for extended transaction definitions

[resolves #104]

Signed-off-by: Mark Paluch <[email protected]>

Author: mp911de

pom.xml

@@ -40,7 +40,7 @@
         <mockito.version>3.5.15</mockito.version>
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
         <project.reporting.outputEncoding>UTF-8</project.reporting.outputEncoding>
-        <r2dbc-spi.version>0.8.1.RELEASE</r2dbc-spi.version>
+        <r2dbc-spi.version>0.9.0.BUILD-SNAPSHOT</r2dbc-spi.version>
         <reactor.version>Dysprosium-SR12</reactor.version>
     </properties>
 

src/main/java/io/r2dbc/pool/PooledConnection.java

@@ -21,6 +21,7 @@
 import io.r2dbc.spi.ConnectionMetadata;
 import io.r2dbc.spi.IsolationLevel;
 import io.r2dbc.spi.Statement;
+import io.r2dbc.spi.TransactionDefinition;
 import io.r2dbc.spi.ValidationDepth;
 import io.r2dbc.spi.Wrapped;
 import org.reactivestreams.Publisher;
@@ -69,7 +70,13 @@ final class PooledConnection implements Connection, Wrapped<Connection> {
     @Override
     public Mono<Void> beginTransaction() {
         assertNotClosed();
-        return Mono.from(this.connection.beginTransaction()).doOnSubscribe(ignore -> this.inTransaction = true);
+        return Mono.from(this.connection.beginTransaction()).doOnSubscribe(ignore -> this.inTransaction = true).doOnError(e -> this.inTransaction = false);
+    }
+
+    @Override
+    public Mono<Void> beginTransaction(TransactionDefinition definition) {
+        assertNotClosed();
+        return Mono.from(this.connection.beginTransaction(definition)).doOnSubscribe(ignore -> this.inTransaction = true).doOnError(e -> this.inTransaction = false);
     }
 
     @Override

src/test/java/io/r2dbc/pool/PooledConnectionUnitTests.java

@@ -17,6 +17,7 @@
 package io.r2dbc.pool;
 
 import io.r2dbc.spi.Connection;
+import io.r2dbc.spi.TransactionDefinition;
 import io.r2dbc.spi.ValidationDepth;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -28,6 +29,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.reset;
@@ -52,6 +54,7 @@ void setUp() {
         when(pooledRefMock.poolable()).thenReturn(connectionMock);
         when(pooledRefMock.release()).thenReturn(Mono.empty());
         when(connectionMock.beginTransaction()).thenReturn(Mono.empty());
+        when(connectionMock.beginTransaction(any())).thenReturn(Mono.empty());
         when(connectionMock.close()).thenReturn(Mono.empty());
         when(connectionMock.validate(ValidationDepth.LOCAL)).thenReturn(Mono.empty());
     }
@@ -71,6 +74,21 @@ void shouldRollbackUnfinishedTransaction() {
         assertThat(wasCalled).isTrue();
     }
 
+    @Test
+    void shouldRollbackUnfinishedExtendedTransaction() {
+
+        AtomicBoolean wasCalled = new AtomicBoolean();
+        when(connectionMock.rollbackTransaction()).thenReturn(Mono.<Void>empty().doOnSuccess(o -> wasCalled.set(true)));
+
+        PooledConnection connection = new PooledConnection(pooledRefMock);
+        connection.beginTransaction(mock(TransactionDefinition.class)).as(StepVerifier::create).verifyComplete();
+
+        connection.close().as(StepVerifier::create).verifyComplete();
+
+        verify(connectionMock).rollbackTransaction();
+        assertThat(wasCalled).isTrue();
+    }
+
     @Test
     void shouldPristineTransactionLeavesTransactionalStateAsIs() {