diff --git a/libsql-jdbc/src/main/java/io/github/magnusp/libsql/jdbc/GeneratedKeys.java b/libsql-jdbc/src/main/java/io/github/magnusp/libsql/jdbc/GeneratedKeys.java new file mode 100644 index 0000000..42ca336 --- /dev/null +++ b/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. + * + *

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 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 statements, List results) + throws SQLException { + clear(); + 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()) { + 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 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(owner, keys); + } + + private static boolean isInsertSql(String sql) { + List script = SqlScript.split(sql); + String statementSql = script.isEmpty() ? sql : script.get(script.size() - 1); + List 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); + } +} diff --git a/libsql-jdbc/src/main/java/io/github/magnusp/libsql/jdbc/LibsqlStatement.java b/libsql-jdbc/src/main/java/io/github/magnusp/libsql/jdbc/LibsqlStatement.java index ad5a10a..ea58b5c 100644 --- a/libsql-jdbc/src/main/java/io/github/magnusp/libsql/jdbc/LibsqlStatement.java +++ b/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 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 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 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); - } + generatedKeys.update(this, sql, result); } } - 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; - generatedKeys = null; - if (keys instanceof LibsqlResultSet libsqlKeys) { - libsqlKeys.closeDetached(); - } else if (keys != null) { - keys.close(); - } + generatedKeys.clear(); } } - private static boolean isInsertSql(String sql) { - List script = SqlScript.split(sql); - String statementSql = script.isEmpty() ? sql : script.get(script.size() - 1); - List 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); } diff --git a/libsql-jdbc/src/test/java/io/github/magnusp/libsql/jdbc/LibsqlGeneratedKeysTest.java b/libsql-jdbc/src/test/java/io/github/magnusp/libsql/jdbc/LibsqlGeneratedKeysTest.java index 1618e9e..8419858 100644 --- a/libsql-jdbc/src/test/java/io/github/magnusp/libsql/jdbc/LibsqlGeneratedKeysTest.java +++ b/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)));