fix(jdbc): match xerial transaction begin and re-begin lifecycle

Closes #40. Part of #54.

setAutoCommit(false) now begins eagerly in the configured transaction mode (DEFERRED by default), and commit/rollback begin the next transaction in the same pipeline. Re-enabling auto-commit commits and closes the Hrana stream.

Author: magnusp

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

@@ -0,0 +1,126 @@
+package com.github.magnusp.libsql.itest;
+
+import com.github.magnusp.libsql.jdbc.LibsqlConnection;
+import org.junit.jupiter.api.Test;
+
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class TransactionLifecycleIT extends SqldContainerBase {
+
+    private String getDriverJdbcUrl() {
+        return "jdbc:libsql:" + getHttpUrl();
+    }
+
+    @Test
+    void deferredTransactionBeginsEagerlyWithoutHoldingAWriteLock() throws Exception {
+        String url = getDriverJdbcUrl();
+        try (Connection lockHolder = DriverManager.getConnection(url);
+             Connection writer = DriverManager.getConnection(url)) {
+            createTable(lockHolder, "deferred_lock");
+
+            lockHolder.setAutoCommit(false);
+
+            ExecutorService executor = Executors.newSingleThreadExecutor();
+            try {
+                Future<Integer> insert = executor.submit(() -> insertRow(writer, "deferred_lock", 1));
+                assertThat(insert.get(5, TimeUnit.SECONDS)).isEqualTo(1);
+            } finally {
+                executor.shutdownNow();
+            }
+        }
+    }
+
+    @Test
+    void immediateTransactionBeginsEagerlyAndHoldsTheWriteLock() throws Exception {
+        String url = getDriverJdbcUrl();
+        try (Connection lockHolder = DriverManager.getConnection(url)) {
+            createTable(lockHolder, "immediate_lock");
+            LibsqlConnection libsql = lockHolder.unwrap(LibsqlConnection.class);
+            libsql.setTransactionMode(LibsqlConnection.TransactionMode.IMMEDIATE);
+
+            lockHolder.setAutoCommit(false);
+
+            assertWriteIsBlockedUntilReleased(url, lockHolder, "immediate_lock", 1);
+        }
+    }
+
+    @Test
+    void commitBeginsTheNextTransactionAndKeepsTheConfiguredLock() throws Exception {
+        String url = getDriverJdbcUrl();
+        try (Connection lockHolder = DriverManager.getConnection(url);
+             Statement statement = lockHolder.createStatement()) {
+            createTable(lockHolder, "commit_rebegin");
+            LibsqlConnection libsql = lockHolder.unwrap(LibsqlConnection.class);
+            libsql.setTransactionMode(LibsqlConnection.TransactionMode.IMMEDIATE);
+            lockHolder.setAutoCommit(false);
+            statement.executeUpdate("INSERT INTO commit_rebegin (id) VALUES (100)");
+
+            lockHolder.commit();
+
+            assertWriteIsBlockedUntilReleased(url, lockHolder, "commit_rebegin", 1);
+        }
+    }
+
+    private static void assertWriteIsBlockedUntilReleased(
+            String url, Connection lockHolder, String table, int id) throws Exception {
+        ExecutorService executor = Executors.newSingleThreadExecutor();
+        try {
+            Future<Object> insert = executor.submit(() -> {
+                try (Connection writer = DriverManager.getConnection(url)) {
+                    return insertRow(writer, table, id);
+                } catch (SQLException e) {
+                    return e;
+                }
+            });
+
+            Object resultWhileLocked;
+            try {
+                resultWhileLocked = insert.get(300, TimeUnit.MILLISECONDS);
+                assertThat(resultWhileLocked)
+                        .as("writer must not succeed while the write lock is held")
+                        .isNotEqualTo(1);
+            } catch (TimeoutException blocked) {
+                resultWhileLocked = null;
+            }
+
+            lockHolder.setAutoCommit(true);
+
+            if (insert.isDone()) {
+                Object completed = insert.get();
+                if (completed instanceof SQLException) {
+                    try (Connection writer = DriverManager.getConnection(url)) {
+                        assertThat(insertRow(writer, table, id)).isEqualTo(1);
+                    }
+                } else {
+                    assertThat(completed).isEqualTo(1);
+                }
+            } else {
+                assertThat(insert.get(20, TimeUnit.SECONDS)).isEqualTo(1);
+            }
+        } finally {
+            executor.shutdownNow();
+        }
+    }
+
+    private static void createTable(Connection connection, String table) throws SQLException {
+        try (Statement statement = connection.createStatement()) {
+            statement.execute("CREATE TABLE " + table + " (id INTEGER PRIMARY KEY)");
+        }
+    }
+
+    private static int insertRow(Connection connection, String table, int id) throws SQLException {
+        try (Statement statement = connection.createStatement()) {
+            return statement.executeUpdate("INSERT INTO " + table + " (id) VALUES (" + id + ")");
+        }
+    }
+}

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

@@ -46,6 +46,35 @@ public class LibsqlConnection implements Connection {
     private boolean readOnly = false;
     private int transactionIsolation = TRANSACTION_SERIALIZABLE;
     private boolean closed = false;
+    private TransactionMode transactionMode = TransactionMode.DEFERRED;
+
+    /**
+     * Transaction begin mode, matching xerial sqlite-jdbc's
+     * {@code SQLiteConfig.TransactionMode}.
+     */
+    public enum TransactionMode {
+        DEFERRED("BEGIN"),
+        IMMEDIATE("BEGIN IMMEDIATE"),
+        EXCLUSIVE("BEGIN EXCLUSIVE");
+
+        private final String beginSql;
+
+        TransactionMode(String beginSql) {
+            this.beginSql = beginSql;
+        }
+
+        String beginSql() {
+            return beginSql;
+        }
+
+        public static TransactionMode from(String mode) {
+            try {
+                return valueOf(mode.toUpperCase());
+            } catch (RuntimeException e) {
+                throw new IllegalArgumentException("Unknown transaction mode: " + mode, e);
+            }
+        }
+    }
 
     // Interactive stream state
     private String baton = null;
@@ -159,15 +188,33 @@ private int[] extractBatchResults(PipelineRespBody resp, int expectedSize) throw
 
     private void ensureTransactionStarted() throws SQLException {
         if (baton == null) {
-            // Mitigation for Bug 2: Read-write transactions start with BEGIN IMMEDIATE
-            String beginSql = readOnly ? "BEGIN DEFERRED" : "BEGIN IMMEDIATE";
-            PipelineReqBody req = new PipelineReqBody(null, List.of(new StreamRequest.Execute(new Stmt(beginSql))));
-            PipelineRespBody resp = client.sendPipeline(req, baseUrl);
-            updateStreamState(resp);
-            if (resp.results().isEmpty() || resp.results().get(0) instanceof StreamResult.Error) {
-                closeStreamSilently();
-                throw new SQLException("Failed to begin transaction");
-            }
+            beginTransaction();
+        }
+    }
+
+    private void beginTransaction() throws SQLException {
+        PipelineRespBody resp;
+        try {
+            PipelineReqBody req = new PipelineReqBody(null,
+                    List.of(new StreamRequest.Execute(new Stmt(transactionMode.beginSql()))));
+            resp = client.sendPipeline(req, baseUrl);
+        } catch (LibsqlException e) {
+            throw LibsqlErrorMapper.toSQLException(e);
+        }
+        updateStreamState(resp);
+        if (resp.results().isEmpty()) {
+            closeStreamSilently();
+            throw new SQLException("Empty result received from server");
+        }
+        StreamResult result = resp.results().get(0);
+        if (result instanceof StreamResult.Error error) {
+            closeStreamSilently();
+            throw new SQLException(error.error().message());
+        }
+        if (!(result instanceof StreamResult.Ok ok)
+                || !(ok.response() instanceof StreamResponse.Execute)) {
+            closeStreamSilently();
+            throw new SQLException("Unexpected response: " + result);
         }
     }
 
@@ -214,15 +261,33 @@ public String nativeSQL(String sql) throws SQLException {
         return sql;
     }
 
+    /**
+     * Matches xerial sqlite-jdbc 3.53.4.0:
+     *
+     * <ul>
+     *   <li>Disabling auto-commit immediately begins a transaction in the configured
+     *       {@link TransactionMode}. The default is {@link TransactionMode#DEFERRED}, so an idle
+     *       non-auto-commit connection holds no lock until its first statement.
+     *   <li>{@link #commit()} and {@link #rollback()} end the current transaction and immediately
+     *       begin the next one, so a non-auto-commit connection always has a transaction open.
+     *   <li>Re-enabling auto-commit commits the open transaction and ends the Hrana transaction
+     *       stream.
+     * </ul>
+     *
+     * <p>Use {@link #setTransactionMode(TransactionMode)} to select {@code IMMEDIATE} or {@code
+     * EXCLUSIVE}, which acquire the write lock when the transaction begins.
+     */
     @Override
     public void setAutoCommit(boolean autoCommit) throws SQLException {
         checkClosed();
         if (this.autoCommit == autoCommit) return;
-        if (!this.autoCommit && autoCommit) {
-            // Committing active transaction
-            commit();
+        if (autoCommit) {
+            endTransaction("COMMIT", false);
+            this.autoCommit = true;
+        } else {
+            this.autoCommit = false;
+            beginTransaction();
         }
-        this.autoCommit = autoCommit;
     }
 
     @Override
@@ -237,7 +302,7 @@ public void commit() throws SQLException {
         if (autoCommit) {
             throw new SQLException("Cannot commit when in autoCommit mode");
         }
-        endTransaction("COMMIT");
+        endTransaction("COMMIT", true);
     }
 
     @Override
@@ -246,18 +311,26 @@ public void rollback() throws SQLException {
         if (autoCommit) {
             throw new SQLException("Cannot rollback when in autoCommit mode");
         }
-        endTransaction("ROLLBACK");
+        endTransaction("ROLLBACK", true);
     }
 
-    private void endTransaction(String sql) throws SQLException {
+    private void endTransaction(String sql, boolean rebegin) throws SQLException {
         if (baton == null) {
+            if (rebegin) {
+                throw new SQLException("No transaction is active");
+            }
             return;
         }
 
+        List<StreamRequest> requests = new ArrayList<>(2);
+        requests.add(new StreamRequest.Execute(new Stmt(sql)));
+        if (rebegin) {
+            requests.add(new StreamRequest.Execute(new Stmt(transactionMode.beginSql())));
+        }
+
         PipelineRespBody resp;
         try {
-            PipelineReqBody req = new PipelineReqBody(baton, List.of(new StreamRequest.Execute(new Stmt(sql))));
-            resp = client.sendPipeline(req, baseUrl);
+            resp = client.sendPipeline(new PipelineReqBody(baton, requests), baseUrl);
         } catch (LibsqlException e) {
             throw LibsqlErrorMapper.toSQLException(e);
         }
@@ -274,10 +347,33 @@ private void endTransaction(String sql) throws SQLException {
             throw new SQLException("Unexpected response: " + result);
         }
 
-        if (resp.baton() != null) {
-            this.baton = resp.baton();
+        if (!rebegin) {
+            updateStreamState(resp);
+            closeStreamSilently();
+            return;
+        }
+
+        if (resp.results().size() < 2) {
+            closeStreamSilently();
+            throw new SQLException("Missing response for transaction begin");
         }
-        closeStreamSilently();
+        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
@@ -319,6 +415,27 @@ public String getCatalog() throws SQLException {
         return null;
     }
 
+    /**
+     * Selects the mode used by the next transaction begin. Matches xerial sqlite-jdbc's {@code
+     * SQLiteConfig.setTransactionMode}; the default is {@link TransactionMode#DEFERRED}.
+     */
+    public void setTransactionMode(TransactionMode mode) throws SQLException {
+        checkClosed();
+        if (mode == null) {
+            throw new IllegalArgumentException("Transaction mode cannot be null");
+        }
+        this.transactionMode = mode;
+    }
+
+    public void setTransactionMode(String mode) throws SQLException {
+        setTransactionMode(TransactionMode.from(mode));
+    }
+
+    public TransactionMode getTransactionMode() throws SQLException {
+        checkClosed();
+        return transactionMode;
+    }
+
     @Override
     public void setTransactionIsolation(int level) throws SQLException {
         checkClosed();