fix(jdbc): reset generated keys and result sets on every execution path

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())) {

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<Stmt> statements) throws SQLException {
+        List<StmtResult> 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<StmtResult> executeBatchResults(List<Stmt> 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<Stmt> statements) throws SQLException {
         }
     }
 
-    private int[] extractBatchResults(PipelineRespBody resp, int expectedSize) throws SQLException {
+    private List<StmtResult> 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<StmtResult> 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<StmtResult> results) {
+        int[] updateCounts = new int[results.size()];
+        for (int i = 0; i < results.size(); i++) {
+            updateCounts[i] = (int) results.get(i).affectedRowCount();
         }
         return updateCounts;
     }

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

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<Stmt> statements, List<StmtResult> results)
+            throws SQLException {
+        synchronized (stateLock) {
+            clearGeneratedKeys();
+            List<Long> 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<Long> rowids) {
+        List<Row> 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<Stmt> statements = new ArrayList<>(batchStatements);
         try {
-            return connection.executeBatch(batchStatements);
+            beginExecution();
+            List<StmtResult> 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();
         }