refactor(jdbc): extract generated-key state

libsql-jdbc/src/main/java/io/github/magnusp/libsql/jdbc/GeneratedKeys.java

@@ -0,0 +1,114 @@
+package io.github.magnusp.libsql.jdbc;
+
+import io.github.magnusp.libsql.hrana.codec.Col;
+import io.github.magnusp.libsql.hrana.codec.Row;
+import io.github.magnusp.libsql.hrana.codec.Stmt;
+import io.github.magnusp.libsql.hrana.codec.StmtResult;
+import io.github.magnusp.libsql.hrana.codec.Value;
+
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+
+/**
+ * Holds the generated-key result of one statement.
+ *
+ * <p>The owning statement guards every method with its {@code stateLock}, the same lock that
+ * guards the current result set. A key result is built from the rowid the server sent on the
+ * statement result, never from a follow-up query. A statement that did not insert reports an empty
+ * one-row-wide result whose column is named {@code 1}, matching sqlite-jdbc 3.53.4.0.
+ */
+final class GeneratedKeys {
+
+    private static final Set<String> STATEMENT_KEYWORDS =
+            Set.of("SELECT", "VALUES", "INSERT", "REPLACE", "UPDATE", "DELETE");
+
+    private ResultSet current;
+
+    void update(LibsqlStatement owner, String sql, StmtResult result) throws SQLException {
+        clear();
+        if (sql != null
+                && result != null
+                && result.lastInsertRowid() != null
+                && isInsertSql(sql)) {
+            current = resultSet(owner, 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.
+     */
+    void updateBatch(LibsqlStatement owner, List<Stmt> statements, List<StmtResult> results)
+            throws SQLException {
+        clear();
+        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()) {
+            current = resultSet(owner, rowids);
+        }
+    }
+
+    ResultSet currentOrEmpty(LibsqlStatement owner) {
+        if (current == null) {
+            StmtResult empty = new StmtResult(List.of(new Col("1", null)), List.of(), 0, null);
+            current = new LibsqlResultSet(owner, empty);
+        }
+        return current;
+    }
+
+    void clear() throws SQLException {
+        ResultSet keys = current;
+        current = null;
+        if (keys instanceof LibsqlResultSet libsqlKeys) {
+            libsqlKeys.closeDetached();
+        } else if (keys != null) {
+            keys.close();
+        }
+    }
+
+    private static LibsqlResultSet resultSet(LibsqlStatement owner, 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(owner, keys);
+    }
+
+    private static boolean isInsertSql(String sql) {
+        List<String> script = SqlScript.split(sql);
+        String statementSql = script.isEmpty() ? sql : script.get(script.size() - 1);
+        List<SqlTokenizer.Token> tokens = SqlTokenizer.tokenize(statementSql);
+        if (tokens == null) {
+            return false;
+        }
+        String first = SqlTokenizer.wordAt(tokens, 0, 0);
+        if ("INSERT".equals(first) || "REPLACE".equals(first)) {
+            return true;
+        }
+        if (!"WITH".equals(first)) {
+            return false;
+        }
+        int index = SqlTokenizer.indexOfAnyWord(tokens, 1, tokens.size(), 0, STATEMENT_KEYWORDS);
+        String statement = SqlTokenizer.wordAt(tokens, index, 0);
+        return "INSERT".equals(statement) || "REPLACE".equals(statement);
+    }
+}

libsql-jdbc/src/main/java/io/github/magnusp/libsql/jdbc/LibsqlStatement.java

@@ -1,7 +1,5 @@
 package io.github.magnusp.libsql.jdbc;
 
-import io.github.magnusp.libsql.hrana.codec.Col;
-import io.github.magnusp.libsql.hrana.codec.Row;
 import io.github.magnusp.libsql.hrana.codec.Stmt;
 import io.github.magnusp.libsql.hrana.codec.StmtResult;
 import io.github.magnusp.libsql.hrana.codec.Value;
@@ -15,7 +13,6 @@
 import java.sql.Statement;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Set;
 
 /**
  * Statement over the Hrana transport.
@@ -74,16 +71,13 @@ public class LibsqlStatement implements Statement {
     protected boolean closed = false;
     protected boolean closeOnCompletion = false;
     protected final List<Stmt> batchStatements = new ArrayList<>();
-    private ResultSet generatedKeys = null;
+    private final GeneratedKeys generatedKeys = new GeneratedKeys();
     private boolean resultSetRequested = false;
     private boolean closing = false;
     private int fetchSize = 0;
     private int maxRows = 0;
     private boolean escapeProcessing = true;
 
-    private static final Set<String> STATEMENT_KEYWORDS =
-            Set.of("SELECT", "VALUES", "INSERT", "REPLACE", "UPDATE", "DELETE");
-
     public LibsqlStatement(LibsqlConnection connection) {
         this.connection = connection;
     }
@@ -342,96 +336,16 @@ protected void closeCurrentResultSet() throws SQLException {
 
     protected void updateGeneratedKeys(String sql, StmtResult result) throws SQLException {
         synchronized (stateLock) {
-            clearGeneratedKeys();
-            if (sql != null
-                    && result != null
-                    && result.lastInsertRowid() != null
-                    && isInsertSql(sql)) {
-                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);
-            }
+            generatedKeys.update(this, sql, result);
         }
     }
 
-    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;
-            generatedKeys = null;
-            if (keys instanceof LibsqlResultSet libsqlKeys) {
-                libsqlKeys.closeDetached();
-            } else if (keys != null) {
-                keys.close();
-            }
+            generatedKeys.clear();
         }
     }
 
-    private static boolean isInsertSql(String sql) {
-        List<String> script = SqlScript.split(sql);
-        String statementSql = script.isEmpty() ? sql : script.get(script.size() - 1);
-        List<SqlTokenizer.Token> tokens = SqlTokenizer.tokenize(statementSql);
-        if (tokens == null) {
-            return false;
-        }
-        String first = SqlTokenizer.wordAt(tokens, 0, 0);
-        if ("INSERT".equals(first) || "REPLACE".equals(first)) {
-            return true;
-        }
-        if (!"WITH".equals(first)) {
-            return false;
-        }
-        int index = SqlTokenizer.indexOfAnyWord(tokens, 1, tokens.size(), 0, STATEMENT_KEYWORDS);
-        String statement = SqlTokenizer.wordAt(tokens, index, 0);
-        return "INSERT".equals(statement) || "REPLACE".equals(statement);
-    }
-
-    private ResultSet emptyGeneratedKeys() {
-        StmtResult empty = new StmtResult(
-                List.of(new Col("1", null)),
-                List.of(),
-                0,
-                null);
-        return new LibsqlResultSet(this, empty);
-    }
-
     @Override
     public ResultSet getResultSet() throws SQLException {
         synchronized (stateLock) {
@@ -514,7 +428,7 @@ public int[] executeBatch() throws SQLException {
                 affectedRows += result.affectedRowCount();
             }
             synchronized (stateLock) {
-                updateBatchGeneratedKeys(statements, results);
+                generatedKeys.updateBatch(this, statements, results);
                 updateCount = affectedRows;
                 resultSetRequested = false;
             }
@@ -653,10 +567,7 @@ public boolean getMoreResults(int current) throws SQLException {
     public ResultSet getGeneratedKeys() throws SQLException {
         synchronized (stateLock) {
             checkClosed();
-            if (generatedKeys == null) {
-                generatedKeys = emptyGeneratedKeys();
-            }
-            return generatedKeys;
+            return generatedKeys.currentOrEmpty(this);
         }
     }
     @Override public int executeUpdate(String sql, int autoGeneratedKeys) throws SQLException { return executeUpdate(sql); }

libsql-jdbc/src/test/java/io/github/magnusp/libsql/jdbc/LibsqlGeneratedKeysTest.java

@@ -410,6 +410,34 @@ void failedBatchLeavesNoGeneratedKeys() throws SQLException {
         }
     }
 
+    @Test
+    void emptyGeneratedKeysUseTheXerialColumnName() throws SQLException {
+        RecordingConnection connection = new RecordingConnection(insertResult(1L));
+
+        try (Statement statement = connection.createStatement();
+                ResultSet keys = statement.getGeneratedKeys()) {
+            assertThat(keys.getMetaData().getColumnName(1)).isEqualTo("1");
+        }
+    }
+
+    @Test
+    void mixedBatchReportsOnlyTheInsertStepKeys() throws SQLException {
+        RecordingConnection connection = new RecordingConnection();
+        connection.batchResults(insertResult(21L), updateResult(1, 99L));
+
+        try (Statement statement = connection.createStatement()) {
+            statement.addBatch("insert into k (v) values ('a')");
+            statement.addBatch("update k set v = 'b'");
+
+            assertThat(statement.executeBatch()).containsExactly(1, 1);
+            try (ResultSet keys = statement.getGeneratedKeys()) {
+                assertThat(keys.next()).isTrue();
+                assertThat(keys.getLong(1)).isEqualTo(21L);
+                assertThat(keys.next()).isFalse();
+            }
+        }
+    }
+
     @Test
     void supportsGetGeneratedKeysMatchesImplementedBehavior() throws SQLException {
         LibsqlDatabaseMetaData metadata = new LibsqlDatabaseMetaData(new RecordingConnection(insertResult(1L)));