Allow access to connection from DbOperations

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - [PR #40](https://github.com/itsallcode/simple-jdbc/pull/40): Verify that operation on `SimpleConnection` are allowed
 - [PR #41](https://github.com/itsallcode/simple-jdbc/pull/41): Allow direct access to `PreparedStatement` for batch insert
 - [PR #42](https://github.com/itsallcode/simple-jdbc/pull/42): Allow direct access to `Connection`
+- [PR #43](https://github.com/itsallcode/simple-jdbc/pull/43): Allow direct access to `Connection` from `DbOperations`
 
 ## [0.9.0] - 2024-12-23
 

src/main/java/org/itsallcode/jdbc/DbOperations.java

@@ -1,5 +1,6 @@
 package org.itsallcode.jdbc;
 
+import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.util.Iterator;
 import java.util.List;
@@ -135,6 +136,15 @@ <T> SimpleResultSet<T> query(final String sql, final PreparedStatementSetter pre
      */
     <T> RowBatchInsertBuilder<T> batchInsert(final Class<T> rowType);
 
+    /**
+     * Get the original wrapped connection.
+     * <p>
+     * Use this in case of missing features in {@link DbOperations}.
+     * 
+     * @return original wrapped connection
+     */
+    public Connection getOriginalConnection();
+
     @Override
     void close();
 }

src/main/java/org/itsallcode/jdbc/SimpleConnection.java

@@ -106,14 +106,8 @@ public <T> RowBatchInsertBuilder<T> batchInsert(final Class<T> rowType) {
         return connection.rowBatchInsert();
     }
 
-    /**
-     * Get the original wrapped connection.
-     * <p>
-     * Use this in case of missing features in {@link SimpleConnection}.
-     * 
-     * @return original wrapped connection
-     */
     public Connection getOriginalConnection() {
+        checkOperationAllowed();
         return connection.getOriginalConnection();
     }
 

src/main/java/org/itsallcode/jdbc/Transaction.java

@@ -111,6 +111,11 @@ public <T> RowBatchInsertBuilder<T> batchInsert(final Class<T> rowType) {
         return connection.rowBatchInsert();
     }
 
+    public Connection getOriginalConnection() {
+        checkOperationAllowed();
+        return connection.getOriginalConnection();
+    }
+
     private void checkOperationAllowed() {
         if (this.closed) {
             throw new IllegalStateException("Operation not allowed on closed transaction");

src/test/java/org/itsallcode/jdbc/SimpleConnectionTest.java

@@ -71,6 +71,7 @@ static Stream<Arguments> operations() {
                 operation(con -> con.query("sql", rowMapperMock)),
                 operation(con -> con.query("sql", preparedStatementSetterMock, rowMapperMock)),
                 operation(con -> con.query("sql", List.of(), rowMapperMock)),
+                operation(con -> con.getOriginalConnection()),
                 operation(con -> con.batchInsert()),
                 operation(con -> con.batchInsert(null)));
     }

src/test/java/org/itsallcode/jdbc/TransactionTest.java

@@ -1,11 +1,13 @@
 package org.itsallcode.jdbc;
 
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.same;
 import static org.mockito.Mockito.*;
 
+import java.sql.Connection;
 import java.util.List;
 import java.util.function.Consumer;
 import java.util.stream.Stream;
@@ -74,6 +76,7 @@ static Stream<Arguments> operations() {
                 operation(tx -> tx.query("sql", List.of(), rowMapperMock)),
                 operation(tx -> tx.batchInsert()),
                 operation(tx -> tx.batchInsert(null)),
+                operation(tx -> tx.getOriginalConnection()),
                 operation(tx -> tx.commit()),
                 operation(tx -> tx.rollback()));
     }
@@ -198,6 +201,13 @@ void rollbackCallsTransactionFinishedCallback() {
         verify(transactionFinishedCallbackMock).accept(same(testee));
     }
 
+    @Test
+    void getOriginalConnection() {
+        final Connection jdbcConnectionMock = mock(Connection.class);
+        when(connectionMock.getOriginalConnection()).thenReturn(jdbcConnectionMock);
+        assertThat(startTransaction().getOriginalConnection()).isSameAs(jdbcConnectionMock);
+    }
+
     private Transaction startTransaction() {
         return Transaction.start(connectionMock, transactionFinishedCallbackMock);
     }