From f0fbdf53fc33648716f19a77ca38728184d9692c Mon Sep 17 00:00:00 2001 From: Magnus Persson <1431685+magnusp@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:22:54 +0200 Subject: [PATCH] fix(jdbc): reset generated keys and result sets on every execution path --- .../magnusp/libsql/itest/GeneratedKeysIT.java | 44 ++++++ .../magnusp/libsql/jdbc/LibsqlConnection.java | 34 +++- .../libsql/jdbc/LibsqlPreparedStatement.java | 4 +- .../magnusp/libsql/jdbc/LibsqlStatement.java | 94 +++++++++-- .../libsql/jdbc/LibsqlGeneratedKeysTest.java | 148 +++++++++++++++++- 5 files changed, 304 insertions(+), 20 deletions(-) diff --git a/libsql-itests/src/test/java/com/github/magnusp/libsql/itest/GeneratedKeysIT.java b/libsql-itests/src/test/java/com/github/magnusp/libsql/itest/GeneratedKeysIT.java index d2edbeb..12b0901 100644 --- a/libsql-itests/src/test/java/com/github/magnusp/libsql/itest/GeneratedKeysIT.java +++ b/libsql-itests/src/test/java/com/github/magnusp/libsql/itest/GeneratedKeysIT.java @@ -80,6 +80,50 @@ void generatedKeysComeFromTheInsertingSession() throws SQLException { } } + @Test + void batchReportsThatBatchesKeysNotAnEarlierStatementsKey() throws SQLException { + try (Connection conn = DriverManager.getConnection(getDriverJdbcUrl())) { + try (Statement s = conn.createStatement()) { + s.execute("create table generated_keys_batch (id integer primary key autoincrement, v text)"); + } + + try (Statement s = conn.createStatement()) { + s.executeUpdate("insert into generated_keys_batch (v) values ('a')"); + s.addBatch("insert into generated_keys_batch (v) values ('b')"); + s.addBatch("insert into generated_keys_batch (v) values ('c')"); + + assertThat(s.executeBatch()).containsExactly(1, 1); + + try (ResultSet keys = s.getGeneratedKeys()) { + assertThat(keys.next()).isTrue(); + assertThat(keys.getLong(1)).isEqualTo(2L); + assertThat(keys.next()).isTrue(); + assertThat(keys.getLong(1)).isEqualTo(3L); + assertThat(keys.next()).isFalse(); + } + } + + try (PreparedStatement ps = conn.prepareStatement( + "insert into generated_keys_batch (v) values (?)", + Statement.RETURN_GENERATED_KEYS)) { + ps.setString(1, "d"); + ps.addBatch(); + ps.setString(1, "e"); + ps.addBatch(); + + assertThat(ps.executeBatch()).containsExactly(1, 1); + + try (ResultSet keys = ps.getGeneratedKeys()) { + assertThat(keys.next()).isTrue(); + assertThat(keys.getLong(1)).isEqualTo(4L); + assertThat(keys.next()).isTrue(); + assertThat(keys.getLong(1)).isEqualTo(5L); + assertThat(keys.next()).isFalse(); + } + } + } + } + @Test void insertInsideTransactionReturnsGeneratedKey() throws SQLException { try (Connection conn = DriverManager.getConnection(getDriverJdbcUrl())) { diff --git a/libsql-jdbc/src/main/java/com/github/magnusp/libsql/jdbc/LibsqlConnection.java b/libsql-jdbc/src/main/java/com/github/magnusp/libsql/jdbc/LibsqlConnection.java index 611ce7a..016113d 100644 --- a/libsql-jdbc/src/main/java/com/github/magnusp/libsql/jdbc/LibsqlConnection.java +++ b/libsql-jdbc/src/main/java/com/github/magnusp/libsql/jdbc/LibsqlConnection.java @@ -35,7 +35,6 @@ import java.sql.Statement; import java.sql.Struct; import java.util.ArrayList; -import java.util.Arrays; import java.util.List; import java.util.Map; import java.util.Properties; @@ -132,9 +131,23 @@ public StmtResult executeStatement(Stmt stmt) throws SQLException { } public int[] executeBatch(List statements) throws SQLException { + List results = executeBatchResults(statements); + int[] updateCounts = new int[results.size()]; + for (int i = 0; i < results.size(); i++) { + updateCounts[i] = (int) results.get(i).affectedRowCount(); + } + return updateCounts; + } + + /** + * Executes a batch and returns the {@code BatchResult} step results, one per statement. The + * caller derives both the update counts and the generated keys from these results, so the + * driver needs no follow-up query and no extra round trip. + */ + List executeBatchResults(List statements) throws SQLException { checkClosed(); if (statements == null || statements.isEmpty()) { - return new int[0]; + return List.of(); } try { @@ -165,7 +178,8 @@ public int[] executeBatch(List statements) throws SQLException { } } - private int[] extractBatchResults(PipelineRespBody resp, int expectedSize) throws SQLException { + private List extractBatchResults(PipelineRespBody resp, int expectedSize) + throws SQLException { if (resp.results().isEmpty()) { throw new SQLException("Empty response for batch execution"); } @@ -178,12 +192,12 @@ private int[] extractBatchResults(PipelineRespBody resp, int expectedSize) throw throw new SQLException("Unexpected batch response: " + response); } - int[] updateCounts = new int[expectedSize]; + List results = new ArrayList<>(expectedSize); for (int i = 0; i < expectedSize; i++) { HranaError error = batchResp.result().stepErrors().get(i); if (error != null) { // JDBC stops the batch at the first failure and reports only the steps that ran. - throw new BatchUpdateException(error.message(), Arrays.copyOf(updateCounts, i)); + throw new BatchUpdateException(error.message(), steppedUpdateCounts(results)); } StmtResult result = batchResp.result().stepResults().get(i); if (result == null) { @@ -192,7 +206,15 @@ private int[] extractBatchResults(PipelineRespBody resp, int expectedSize) throw throw new SQLException( "Protocol error: batch step " + i + " has neither a result nor an error"); } - updateCounts[i] = (int) result.affectedRowCount(); + results.add(result); + } + return results; + } + + private static int[] steppedUpdateCounts(List results) { + int[] updateCounts = new int[results.size()]; + for (int i = 0; i < results.size(); i++) { + updateCounts[i] = (int) results.get(i).affectedRowCount(); } return updateCounts; } diff --git a/libsql-jdbc/src/main/java/com/github/magnusp/libsql/jdbc/LibsqlPreparedStatement.java b/libsql-jdbc/src/main/java/com/github/magnusp/libsql/jdbc/LibsqlPreparedStatement.java index 535a9c4..8f3ecf8 100644 --- a/libsql-jdbc/src/main/java/com/github/magnusp/libsql/jdbc/LibsqlPreparedStatement.java +++ b/libsql-jdbc/src/main/java/com/github/magnusp/libsql/jdbc/LibsqlPreparedStatement.java @@ -65,7 +65,7 @@ public ResultSet executeQuery() throws SQLException { throw new SQLException("Query does not return results"); } LibsqlResultSet resultSet = new LibsqlResultSet(this, result, sql, false); - assignResultSet(resultSet, true); + assignResultSet(resultSet, result, sql, true); return resultSet; } @@ -222,7 +222,7 @@ public boolean execute() throws SQLException { private StmtResult executeCurrent() throws SQLException { checkClosed(); - closeCurrentResultSet(); + beginExecution(); checkClosed(); return connection.executeStatement(buildCurrentStmt()); } diff --git a/libsql-jdbc/src/main/java/com/github/magnusp/libsql/jdbc/LibsqlStatement.java b/libsql-jdbc/src/main/java/com/github/magnusp/libsql/jdbc/LibsqlStatement.java index 877cc7c..1fc946b 100644 --- a/libsql-jdbc/src/main/java/com/github/magnusp/libsql/jdbc/LibsqlStatement.java +++ b/libsql-jdbc/src/main/java/com/github/magnusp/libsql/jdbc/LibsqlStatement.java @@ -86,12 +86,18 @@ public ResultSet executeQuery(String sql) throws SQLException { throw new SQLException("Query did not produce a ResultSet: " + sql); } LibsqlResultSet resultSet = new LibsqlResultSet(this, result, sql, false); - assignResultSet(resultSet, true); + assignResultSet(resultSet, result, sql, true); return resultSet; } protected void assignResultSet(LibsqlResultSet resultSet, boolean requested) throws SQLException { + assignResultSet(resultSet, null, null, requested); + } + + protected void assignResultSet( + LibsqlResultSet resultSet, StmtResult result, String sql, boolean requested) + throws SQLException { if (!isClosed()) { beforeResultSetAssignment(); } @@ -103,6 +109,7 @@ protected void assignResultSet(LibsqlResultSet resultSet, boolean requested) updateCount = -1; return; } + updateGeneratedKeys(sql, result); currentResultSet = resultSet; updateCount = -1; resultSetRequested = requested; @@ -149,10 +156,23 @@ protected boolean executeResult(String sql, StmtResult result) throws SQLExcepti } private StmtResult executeSql(String sql) throws SQLException { - closeCurrentResultSet(); + beginExecution(); return connection.executeStatement(new Stmt(sql)); } + /** + * Resets the state every execution path starts from: the previous result set is closed and the + * previous generated keys are discarded. Without this, a path that does not produce keys of + * its own would report the keys of the statement before it. + */ + protected void beginExecution() throws SQLException { + synchronized (stateLock) { + closeCurrentResultSet(); + clearGeneratedKeys(); + updateCount = -1; + } + } + protected void closeCurrentResultSet() throws SQLException { synchronized (stateLock) { LibsqlResultSet resultSet = currentResultSet; @@ -186,17 +206,51 @@ protected void updateGeneratedKeys(String sql, StmtResult result) throws SQLExce && result != null && result.lastInsertRowid() != null && isInsertSql(sql)) { - long rowid = result.lastInsertRowid(); - StmtResult keys = new StmtResult( - List.of(new Col("last_insert_rowid()", null)), - List.of(new Row(List.of(Value.of(rowid)))), - 0, - rowid); - generatedKeys = new LibsqlResultSet(this, keys); + generatedKeys = generatedKeysResultSet(List.of(result.lastInsertRowid())); } } } + /** + * Builds the generated keys for a batch from the {@code BatchResult} step results. Each insert + * step contributes one row, matching the java.sql shape of a result set with one row per + * generated key. Steps that are not inserts contribute nothing because Hrana leaves {@code + * last_insert_rowid} undefined for them. + */ + private void updateBatchGeneratedKeys(List statements, List results) + throws SQLException { + synchronized (stateLock) { + clearGeneratedKeys(); + List rowids = new ArrayList<>(); + for (int i = 0; i < results.size(); i++) { + StmtResult result = results.get(i); + String sql = statements.get(i).sql(); + if (result != null + && result.lastInsertRowid() != null + && sql != null + && isInsertSql(sql)) { + rowids.add(result.lastInsertRowid()); + } + } + if (!rowids.isEmpty()) { + generatedKeys = generatedKeysResultSet(rowids); + } + } + } + + private LibsqlResultSet generatedKeysResultSet(List rowids) { + List rows = new ArrayList<>(rowids.size()); + for (Long rowid : rowids) { + rows.add(new Row(List.of(Value.of(rowid)))); + } + StmtResult keys = new StmtResult( + List.of(new Col("last_insert_rowid()", null)), + rows, + 0, + rowids.get(rowids.size() - 1)); + return new LibsqlResultSet(this, keys); + } + protected void clearGeneratedKeys() throws SQLException { synchronized (stateLock) { ResultSet keys = generatedKeys; @@ -357,11 +411,31 @@ public void clearBatch() throws SQLException { batchStatements.clear(); } + /** + * Executes the batch and reports one update count per step. The generated keys and the update + * count are rebuilt from the step results, so neither can leak from an earlier execution. The + * update count is the total affected rows across the batch; JDBC leaves {@code getUpdateCount} + * unspecified after a batch, and xerial leaves the previous value in place. + */ @Override public int[] executeBatch() throws SQLException { checkClosed(); + List statements = new ArrayList<>(batchStatements); try { - return connection.executeBatch(batchStatements); + beginExecution(); + List results = connection.executeBatchResults(statements); + int[] updateCounts = new int[results.size()]; + long affectedRows = 0; + for (int i = 0; i < results.size(); i++) { + updateCounts[i] = (int) results.get(i).affectedRowCount(); + affectedRows += results.get(i).affectedRowCount(); + } + synchronized (stateLock) { + updateBatchGeneratedKeys(statements, results); + updateCount = affectedRows; + resultSetRequested = false; + } + return updateCounts; } finally { clearBatch(); } diff --git a/libsql-jdbc/src/test/java/com/github/magnusp/libsql/jdbc/LibsqlGeneratedKeysTest.java b/libsql-jdbc/src/test/java/com/github/magnusp/libsql/jdbc/LibsqlGeneratedKeysTest.java index fed79df..d20103d 100644 --- a/libsql-jdbc/src/test/java/com/github/magnusp/libsql/jdbc/LibsqlGeneratedKeysTest.java +++ b/libsql-jdbc/src/test/java/com/github/magnusp/libsql/jdbc/LibsqlGeneratedKeysTest.java @@ -7,6 +7,7 @@ import com.github.magnusp.libsql.hrana.codec.Value; import org.junit.jupiter.api.Test; +import java.sql.BatchUpdateException; import java.sql.ResultSet; import java.sql.SQLException; import java.sql.Statement; @@ -15,6 +16,7 @@ import java.util.List; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; class LibsqlGeneratedKeysTest { @@ -168,7 +170,7 @@ void laterNonInsertClearsPreviousGeneratedKey() throws SQLException { } @Test - void executeQueryDoesNotClearPreviousGeneratedKey() throws SQLException { + void executeQuerySelectClearsPreviousGeneratedKey() throws SQLException { StmtResult rows = new StmtResult( List.of(new Col("id", "INTEGER")), List.of(new Row(List.of(Value.of(3L)))), @@ -182,9 +184,127 @@ void executeQueryDoesNotClearPreviousGeneratedKey() throws SQLException { assertThat(resultSet.next()).isTrue(); } + try (ResultSet keys = statement.getGeneratedKeys()) { + assertThat(keys.next()).isFalse(); + } + } + } + + @Test + void executeQueryReturningReportsThisStatementsKeyNotThePreviousOne() throws SQLException { + StmtResult returning = new StmtResult( + List.of(new Col("id", "INTEGER")), + List.of(new Row(List.of(Value.of(9L)))), + 1, + 9L); + RecordingConnection connection = new RecordingConnection(insertResult(3L), returning); + + try (Statement statement = connection.createStatement()) { + statement.executeUpdate("insert into k (v) values ('a')"); + try (ResultSet resultSet = + statement.executeQuery("insert into k (v) values ('b') returning id")) { + assertThat(resultSet.next()).isTrue(); + assertThat(resultSet.getLong(1)).isEqualTo(9L); + } + + try (ResultSet keys = statement.getGeneratedKeys()) { + assertThat(keys.next()).isTrue(); + assertThat(keys.getLong(1)).isEqualTo(9L); + assertThat(keys.next()).isFalse(); + } + } + } + + @Test + void executeThenBatchReportsTheBatchKeysAndUpdateCount() throws SQLException { + RecordingConnection connection = new RecordingConnection(insertResult(5L)); + connection.batchResults(insertResult(6L), insertResult(7L)); + + try (Statement statement = connection.createStatement()) { + statement.executeUpdate("insert into k (v) values ('a')"); + statement.addBatch("insert into k (v) values ('b')"); + statement.addBatch("insert into k (v) values ('c')"); + + assertThat(statement.executeBatch()).containsExactly(1, 1); + assertThat(statement.getUpdateCount()).isEqualTo(2); + + try (ResultSet keys = statement.getGeneratedKeys()) { + assertThat(keys.getMetaData().getColumnCount()).isEqualTo(1); + assertThat(keys.getMetaData().getColumnName(1)).isEqualTo("last_insert_rowid()"); + assertThat(keys.next()).isTrue(); + assertThat(keys.getLong(1)).isEqualTo(6L); + assertThat(keys.next()).isTrue(); + assertThat(keys.getLong(1)).isEqualTo(7L); + assertThat(keys.next()).isFalse(); + } + } + } + + @Test + void returnGeneratedKeysBatchReportsEveryStepKey() throws SQLException { + RecordingConnection connection = new RecordingConnection(); + connection.batchResults(insertResult(11L), insertResult(12L)); + + try (LibsqlPreparedStatement statement = (LibsqlPreparedStatement) connection.prepareStatement( + "insert into k (v) values (?)", Statement.RETURN_GENERATED_KEYS)) { + statement.setString(1, "a"); + statement.addBatch(); + statement.setString(1, "b"); + statement.addBatch(); + + assertThat(statement.executeBatch()).containsExactly(1, 1); + assertThat(statement.getUpdateCount()).isEqualTo(2); + try (ResultSet keys = statement.getGeneratedKeys()) { assertThat(keys.next()).isTrue(); - assertThat(keys.getLong(1)).isEqualTo(3L); + assertThat(keys.getLong(1)).isEqualTo(11L); + assertThat(keys.next()).isTrue(); + assertThat(keys.getLong(1)).isEqualTo(12L); + assertThat(keys.next()).isFalse(); + } + } + } + + @Test + void executeBatchClosesThePreviousResultSet() throws SQLException { + StmtResult rows = new StmtResult( + List.of(new Col("id", "INTEGER")), + List.of(new Row(List.of(Value.of(1L)))), + 0, + null); + RecordingConnection connection = new RecordingConnection(rows); + connection.batchResults(insertResult(6L)); + + try (Statement statement = connection.createStatement()) { + ResultSet previous = statement.executeQuery("select id from k"); + assertThat(previous.isClosed()).isFalse(); + + statement.addBatch("insert into k (v) values ('b')"); + statement.executeBatch(); + + assertThat(previous.isClosed()).isTrue(); + } + } + + @Test + void failedBatchLeavesNoGeneratedKeys() throws SQLException { + RecordingConnection connection = new RecordingConnection(insertResult(5L)); + connection.failingBatch(insertResult(6L), "batch entry 1: constraint failed"); + + try (Statement statement = connection.createStatement()) { + statement.executeUpdate("insert into k (v) values ('a')"); + statement.addBatch("insert into k (v) values ('b')"); + statement.addBatch("insert into k (v) values ('c')"); + + assertThatThrownBy(statement::executeBatch) + .isInstanceOf(BatchUpdateException.class) + .hasMessageContaining("constraint failed") + .satisfies(thrown -> assertThat( + ((BatchUpdateException) thrown).getUpdateCounts()) + .containsExactly(1)); + + try (ResultSet keys = statement.getGeneratedKeys()) { + assertThat(keys.next()).isFalse(); } } } @@ -198,15 +318,39 @@ void supportsGetGeneratedKeysMatchesImplementedBehavior() throws SQLException { private static final class RecordingConnection extends LibsqlConnection { private final Deque results = new ArrayDeque<>(); + private final Deque> batchResults = new ArrayDeque<>(); + private String batchFailure; private RecordingConnection(StmtResult... results) { super(null, null); this.results.addAll(List.of(results)); } + private void batchResults(StmtResult... results) { + batchResults.add(List.of(results)); + } + + private void failingBatch(StmtResult completed, String message) { + batchResults.add(List.of(completed)); + batchFailure = message; + } + @Override public StmtResult executeStatement(Stmt stmt) { return results.size() == 1 ? results.peek() : results.remove(); } + + @Override + public List executeBatchResults(List statements) throws SQLException { + if (batchFailure != null) { + String message = batchFailure; + batchFailure = null; + List completed = batchResults.remove(); + throw new BatchUpdateException( + message, + completed.stream().mapToInt(r -> (int) r.affectedRowCount()).toArray()); + } + return batchResults.remove(); + } } }