fix(jdbc): rotate the baton before checking endTransaction errors (#98)

libsql-itests/src/test/java/com/github/magnusp/libsql/itest/TransactionLifecycleIT.java

@@ -5,6 +5,7 @@
 
 import java.sql.Connection;
 import java.sql.DriverManager;
+import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.concurrent.ExecutorService;
@@ -14,6 +15,7 @@
 import java.util.concurrent.TimeoutException;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 public class TransactionLifecycleIT extends SqldContainerBase {
 
@@ -71,6 +73,40 @@ void commitBeginsTheNextTransactionAndKeepsTheConfiguredLock() throws Exception
         }
     }
 
+    @Test
+    void failedCommitCanBeRolledBackAndTheConnectionStaysUsable() throws Exception {
+        String url = getDriverJdbcUrl();
+        try (Connection connection = DriverManager.getConnection(url)) {
+            connection.setAutoCommit(false);
+            try (Statement statement = connection.createStatement()) {
+                statement.execute("PRAGMA foreign_keys = ON");
+                statement.execute("CREATE TABLE p (id INTEGER PRIMARY KEY)");
+                statement.execute("CREATE TABLE c (pid INTEGER REFERENCES p(id)"
+                        + " DEFERRABLE INITIALLY DEFERRED)");
+                statement.executeUpdate("INSERT INTO c (pid) VALUES (1)");
+            }
+
+            assertThatThrownBy(connection::commit)
+                    .isInstanceOf(SQLException.class)
+                    .hasMessageContaining("FOREIGN KEY");
+
+            connection.rollback();
+
+            try (Statement statement = connection.createStatement()) {
+                statement.execute("CREATE TABLE usable (id INTEGER PRIMARY KEY)");
+                assertThat(statement.executeUpdate("INSERT INTO usable (id) VALUES (1)"))
+                        .isEqualTo(1);
+            }
+            connection.commit();
+
+            try (Statement statement = connection.createStatement();
+                 ResultSet rows = statement.executeQuery("SELECT COUNT(*) FROM usable")) {
+                assertThat(rows.next()).isTrue();
+                assertThat(rows.getInt(1)).isEqualTo(1);
+            }
+        }
+    }
+
     private static void assertWriteIsBlockedUntilReleased(
             String url, Connection lockHolder, String table, int id) throws Exception {
         ExecutorService executor = Executors.newSingleThreadExecutor();

libsql-jdbc/src/main/java/com/github/magnusp/libsql/jdbc/LibsqlConnection.java

@@ -335,6 +335,7 @@ private void endTransaction(String sql, boolean rebegin) throws SQLException {
         } catch (LibsqlException e) {
             throw LibsqlErrorMapper.toSQLException(e);
         }
+        updateStreamState(resp);
 
         if (resp.results().isEmpty()) {
             throw new SQLException("Empty result received from server");
@@ -349,7 +350,6 @@ private void endTransaction(String sql, boolean rebegin) throws SQLException {
         }
 
         if (!rebegin) {
-            updateStreamState(resp);
             closeStreamSilently();
             return;
         }
@@ -360,21 +360,14 @@ private void endTransaction(String sql, boolean rebegin) throws SQLException {
         }
         StreamResult begin = resp.results().get(1);
         if (begin instanceof StreamResult.Error error) {
-            if (resp.baton() != null) {
-                this.baton = resp.baton();
-            }
             closeStreamSilently();
             throw new SQLException(error.error().message());
         }
         if (!(begin instanceof StreamResult.Ok beginOk)
                 || !(beginOk.response() instanceof StreamResponse.Execute)) {
-            if (resp.baton() != null) {
-                this.baton = resp.baton();
-            }
             closeStreamSilently();
             throw new SQLException("Unexpected response: " + begin);
         }
-        updateStreamState(resp);
     }
 
     @Override

libsql-jdbc/src/test/java/com/github/magnusp/libsql/jdbc/LibsqlConnectionTest.java

@@ -123,6 +123,36 @@ void enablingAutoCommitCommitsAndClosesTheOpenTransaction() throws SQLException
 
         assertThat(client.executedStatements()).containsExactly("BEGIN", "COMMIT");
         assertThat(client.requests.get(1).baton()).isEqualTo(BATON);
+        assertThat(client.requests).hasSize(3);
+        assertThat(client.requests.get(2).requests()).containsExactly(StreamRequest.Close.INSTANCE);
+    }
+
+    @Test
+    void failedCommitCanBeRolledBackAndTheConnectionStaysUsable() throws SQLException {
+        AtomicInteger commits = new AtomicInteger();
+        ScriptedClient client = new ScriptedClient((request, index) -> {
+            if ("COMMIT".equals(executedSql(request)) && commits.incrementAndGet() == 1) {
+                return errorResponse(new HranaError(SERVER_ERROR, "SQLITE_CONSTRAINT_FOREIGNKEY"));
+            }
+            return successResponse(request);
+        });
+        LibsqlConnection connection = connectionInTransaction(client);
+
+        assertThatThrownBy(connection::commit)
+                .isInstanceOf(SQLException.class)
+                .hasMessageContaining(SERVER_ERROR);
+
+        connection.rollback();
+        connection.executeStatement(new Stmt("INSERT INTO sample VALUES (2)"));
+
+        assertThat(client.executedStatements()).containsExactly(
+                "BEGIN",
+                "INSERT INTO sample VALUES (1)",
+                "COMMIT",
+                "BEGIN",
+                "ROLLBACK",
+                "BEGIN",
+                "INSERT INTO sample VALUES (2)");
     }
 
     @ParameterizedTest
@@ -131,8 +161,7 @@ void serverErrorIsPropagatedAndTransactionStateIsKept(String command) throws SQL
         AtomicInteger attempts = new AtomicInteger();
         ScriptedClient client = new ScriptedClient((request, index) -> {
             if (command.equals(executedSql(request)) && attempts.incrementAndGet() == 1) {
-                return new PipelineRespBody(request.baton(), null,
-                        List.of(new StreamResult.Error(new HranaError(SERVER_ERROR, "SQLITE_ERROR"))));
+                return errorResponse(new HranaError(SERVER_ERROR, "SQLITE_ERROR"));
             }
             return successResponse(request);
         });
@@ -174,7 +203,7 @@ private static void assertCommandRetriedOnSameStream(ScriptedClient client, Stri
                 .toList();
         assertThat(commandRequests).hasSize(2);
         assertThat(commandRequests).allSatisfy(
-                request -> assertThat(request.baton()).isEqualTo(BATON));
+                request -> assertThat(request.baton()).isNotNull());
     }
 
     private static LibsqlConnection connectionInTransaction(ScriptedClient client) throws SQLException {
@@ -202,7 +231,11 @@ private static PipelineRespBody successResponse(PipelineReqBody request) {
                         new StreamResponse.Execute(new StmtResult(List.of(), List.of(), 0, null))));
             }
         }
-        return new PipelineRespBody(BATON, null, results);
+        return new PipelineRespBody(null, null, results);
+    }
+
+    private static PipelineRespBody errorResponse(HranaError error) {
+        return new PipelineRespBody(null, null, List.of(new StreamResult.Error(error)));
     }
 
     private static String executedSql(PipelineReqBody request) {
@@ -222,6 +255,8 @@ private interface Handler {
 
         private final Handler handler;
         private final List<PipelineReqBody> requests = new ArrayList<>();
+        private String issuedBaton = null;
+        private int batonCount = 0;
 
         private ScriptedClient(Handler handler) {
             super(CONFIG);
@@ -231,7 +266,12 @@ private ScriptedClient(Handler handler) {
         @Override
         public PipelineRespBody sendPipeline(PipelineReqBody request, URI targetUri) {
             requests.add(request);
-            return handler.handle(request, requests.size() - 1);
+            if (request.baton() != null && !request.baton().equals(issuedBaton)) {
+                throw new LibsqlException("stale baton " + request.baton(), "STALE_BATON", 400);
+            }
+            PipelineRespBody response = handler.handle(request, requests.size() - 1);
+            issuedBaton = "baton-" + (++batonCount);
+            return new PipelineRespBody(issuedBaton, response.baseUrl(), response.results());
         }
 
         private List<String> executedStatements() {