diff --git a/core/src/main/java/com/scalar/db/storage/jdbc/JdbcDatabase.java b/core/src/main/java/com/scalar/db/storage/jdbc/JdbcDatabase.java index 92c8cd60d0..ba0146b246 100644 --- a/core/src/main/java/com/scalar/db/storage/jdbc/JdbcDatabase.java +++ b/core/src/main/java/com/scalar/db/storage/jdbc/JdbcDatabase.java @@ -104,6 +104,17 @@ public Scanner scan(Scan scan) throws ExecutionException { close(connection); throw new ExecutionException( CoreError.JDBC_ERROR_OCCURRED_IN_SELECTION.buildMessage(e.getMessage()), e); + } catch (Exception e) { + try { + if (connection != null) { + connection.rollback(); + } + } catch (SQLException ex) { + e.addSuppressed(ex); + } + + close(connection); + throw e; } } @@ -173,6 +184,9 @@ public void mutate(List mutations) throws ExecutionException close(connection); throw new ExecutionException( CoreError.JDBC_ERROR_OCCURRED_IN_MUTATION.buildMessage(e.getMessage()), e); + } catch (Exception e) { + close(connection); + throw e; } try { diff --git a/core/src/test/java/com/scalar/db/storage/jdbc/JdbcDatabaseTest.java b/core/src/test/java/com/scalar/db/storage/jdbc/JdbcDatabaseTest.java index 8b88e36814..27d31879b5 100644 --- a/core/src/test/java/com/scalar/db/storage/jdbc/JdbcDatabaseTest.java +++ b/core/src/test/java/com/scalar/db/storage/jdbc/JdbcDatabaseTest.java @@ -2,6 +2,9 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -125,6 +128,25 @@ public void whenScanOperationExecutedAndScannerClosed_shouldCallJdbcService() th verify(connection).close(); } + @Test + public void + whenScanOperationExecutedAndJdbcServiceThrowsIllegalArgumentException_shouldCloseConnectionAndThrowIllegalArgumentException() + throws Exception { + // Arrange + Exception cause = new IllegalArgumentException("Table not found"); + // Simulate the table not found scenario. + when(jdbcService.getScanner(any(), any())).thenThrow(cause); + + // Act Assert + assertThatThrownBy( + () -> { + Scan scan = new Scan(new Key("p1", "val")).forNamespace(NAMESPACE).forTable(TABLE); + jdbcDatabase.scan(scan); + }) + .isInstanceOf(IllegalArgumentException.class); + verify(connection).close(); + } + @Test public void whenPutOperationExecuted_shouldCallJdbcService() throws Exception { // Arrange @@ -330,4 +352,30 @@ public void mutate_withConflictError_shouldThrowRetriableExecutionException() .isInstanceOf(RetriableExecutionException.class); verify(connection).close(); } + + @Test + public void mutate_WhenSettingAutoCommitFails_ShouldThrowExceptionAndCloseConnection() + throws SQLException, ExecutionException { + // Arrange + Exception exception = new RuntimeException("Failed to set auto-commit"); + doThrow(exception).when(connection).setAutoCommit(anyBoolean()); + + // Act Assert + assertThatThrownBy( + () -> { + Put put = + new Put(new Key("p1", "val1")) + .withValue("v1", "val2") + .forNamespace(NAMESPACE) + .forTable(TABLE); + Delete delete = + new Delete(new Key("p1", "val1")).forNamespace(NAMESPACE).forTable(TABLE); + jdbcDatabase.mutate(Arrays.asList(put, delete)); + }) + .isEqualTo(exception); + verify(connection).setAutoCommit(false); + verify(jdbcService, never()).mutate(any(), any()); + verify(connection, never()).rollback(); + verify(connection).close(); + } }