Fix potential connection leak when using jdbc storage and Scan operation fails because the target table doesn't exist (#2766)

Co-authored-by: Hiroyuki Yamada <[email protected]>

Author: komamitsu

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<? extends Mutation> 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 {

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();
+  }
 }