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 new file mode 100644 index 0000000..d2edbeb --- /dev/null +++ b/libsql-itests/src/test/java/com/github/magnusp/libsql/itest/GeneratedKeysIT.java @@ -0,0 +1,108 @@ +package com.github.magnusp.libsql.itest; + +import org.junit.jupiter.api.Test; + +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; + +import static org.assertj.core.api.Assertions.assertThat; + +public class GeneratedKeysIT extends SqldContainerBase { + + private String getDriverJdbcUrl() { + return "jdbc:libsql:" + getHttpUrl(); + } + + @Test + void statementAndPreparedStatementInsertsReturnTheirGeneratedKeys() throws SQLException { + try (Connection conn = DriverManager.getConnection(getDriverJdbcUrl())) { + try (Statement s = conn.createStatement()) { + s.execute("create table generated_keys_statement (id integer primary key autoincrement, v text)"); + } + + try (Statement s = conn.createStatement()) { + s.executeUpdate("insert into generated_keys_statement (v) values ('a')"); + try (ResultSet keys = s.getGeneratedKeys()) { + assertThat(keys.getMetaData().getColumnCount()).isEqualTo(1); + assertThat(keys.next()).isTrue(); + assertThat(keys.getLong(1)).isEqualTo(1L); + assertThat(keys.next()).isFalse(); + } + + s.executeUpdate("update generated_keys_statement set v = 'b'"); + try (ResultSet keys = s.getGeneratedKeys()) { + assertThat(keys.getMetaData().getColumnCount()).isEqualTo(1); + assertThat(keys.next()).isFalse(); + } + } + + try (PreparedStatement ps = conn.prepareStatement("insert into generated_keys_statement (v) values (?)")) { + ps.setString(1, "c"); + ps.executeUpdate(); + try (ResultSet keys = ps.getGeneratedKeys()) { + assertThat(keys.getMetaData().getColumnCount()).isEqualTo(1); + assertThat(keys.next()).isTrue(); + assertThat(keys.getLong(1)).isEqualTo(2L); + assertThat(keys.next()).isFalse(); + } + } + + assertThat(conn.getMetaData().supportsGetGeneratedKeys()).isTrue(); + } + } + + @Test + void generatedKeysComeFromTheInsertingSession() throws SQLException { + try (Connection first = DriverManager.getConnection(getDriverJdbcUrl()); + Connection second = DriverManager.getConnection(getDriverJdbcUrl())) { + try (Statement s = first.createStatement()) { + s.execute("create table generated_keys_session (id integer primary key autoincrement, v text)"); + } + + try (Statement firstStatement = first.createStatement(); + Statement secondStatement = second.createStatement()) { + firstStatement.executeUpdate("insert into generated_keys_session (v) values ('a')"); + secondStatement.executeUpdate("insert into generated_keys_session (v) values ('b')"); + + try (ResultSet keys = secondStatement.getGeneratedKeys()) { + assertThat(keys.next()).isTrue(); + assertThat(keys.getLong(1)).isEqualTo(2L); + } + try (ResultSet keys = firstStatement.getGeneratedKeys()) { + assertThat(keys.next()).isTrue(); + assertThat(keys.getLong(1)).isEqualTo(1L); + } + } + } + } + + @Test + void insertInsideTransactionReturnsGeneratedKey() throws SQLException { + try (Connection conn = DriverManager.getConnection(getDriverJdbcUrl())) { + try (Statement s = conn.createStatement()) { + s.execute("create table generated_keys_tx (id integer primary key autoincrement, v text)"); + } + + conn.setAutoCommit(false); + try (PreparedStatement ps = conn.prepareStatement("insert into generated_keys_tx (v) values (?)")) { + ps.setString(1, "a"); + ps.executeUpdate(); + try (ResultSet keys = ps.getGeneratedKeys()) { + assertThat(keys.next()).isTrue(); + assertThat(keys.getLong(1)).isEqualTo(1L); + } + } + conn.commit(); + + try (Statement s = conn.createStatement(); + ResultSet rs = s.executeQuery("select count(*) from generated_keys_tx")) { + assertThat(rs.next()).isTrue(); + assertThat(rs.getLong(1)).isEqualTo(1L); + } + } + } +} diff --git a/libsql-jdbc/src/main/java/com/github/magnusp/libsql/jdbc/LibsqlDatabaseMetaData.java b/libsql-jdbc/src/main/java/com/github/magnusp/libsql/jdbc/LibsqlDatabaseMetaData.java index ede57eb..a4fd91c 100644 --- a/libsql-jdbc/src/main/java/com/github/magnusp/libsql/jdbc/LibsqlDatabaseMetaData.java +++ b/libsql-jdbc/src/main/java/com/github/magnusp/libsql/jdbc/LibsqlDatabaseMetaData.java @@ -401,7 +401,7 @@ public Connection getConnection() throws SQLException { @Override public boolean supportsSavepoints() throws SQLException { return false; } @Override public boolean supportsNamedParameters() throws SQLException { return false; } @Override public boolean supportsMultipleOpenResults() throws SQLException { return false; } - @Override public boolean supportsGetGeneratedKeys() throws SQLException { return false; } + @Override public boolean supportsGetGeneratedKeys() throws SQLException { return true; } @Override public ResultSet getSuperTypes(String catalog, String schemaPattern, String typeNamePattern) throws SQLException { throw new UnsupportedOperationException(); } @Override public ResultSet getSuperTables(String catalog, String schemaPattern, String tableNamePattern) throws SQLException { throw new UnsupportedOperationException(); } @Override public ResultSet getAttributes(String catalog, String schemaPattern, String typeNamePattern, String attributeNamePattern) throws SQLException { throw new UnsupportedOperationException(); } 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 7c27397..049ff6e 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 @@ -60,10 +60,12 @@ public ResultSet executeQuery() throws SQLException { if (!returnsResultSet) { throw new SQLException("Query does not return results"); } - execute(); - if (currentResultSet == null) { + StmtResult result = executeCurrent(); + if (result.cols() == null || result.cols().isEmpty()) { throw new SQLException("Query does not return results"); } + this.currentResultSet = new LibsqlResultSet(this, result, sql, false); + this.updateCount = -1; return currentResultSet; } @@ -215,10 +217,8 @@ private record Token(String word, int depth) {} @Override public boolean execute() throws SQLException { - checkClosed(); - closeCurrentResultSet(); - checkClosed(); - StmtResult result = connection.executeStatement(buildCurrentStmt()); + StmtResult result = executeCurrent(); + updateGeneratedKeys(sql, result); if (result.cols() != null && !result.cols().isEmpty()) { this.currentResultSet = new LibsqlResultSet(this, result, sql, false); this.updateCount = -1; @@ -230,6 +230,13 @@ public boolean execute() throws SQLException { } } + private StmtResult executeCurrent() throws SQLException { + checkClosed(); + closeCurrentResultSet(); + checkClosed(); + return connection.executeStatement(buildCurrentStmt()); + } + @Override public void addBatch() throws SQLException { checkClosed(); 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 951f5bc..5393cef 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 @@ -1,7 +1,10 @@ package com.github.magnusp.libsql.jdbc; +import com.github.magnusp.libsql.hrana.codec.Col; +import com.github.magnusp.libsql.hrana.codec.Row; import com.github.magnusp.libsql.hrana.codec.Stmt; import com.github.magnusp.libsql.hrana.codec.StmtResult; +import com.github.magnusp.libsql.hrana.codec.Value; import java.sql.Connection; import java.sql.ResultSet; @@ -19,6 +22,7 @@ public class LibsqlStatement implements Statement { protected boolean closed = false; protected boolean closeOnCompletion = false; protected final List batchStatements = new ArrayList<>(); + private ResultSet generatedKeys = null; public LibsqlStatement(LibsqlConnection connection) { this.connection = connection; @@ -33,10 +37,12 @@ protected void checkClosed() throws SQLException { @Override public ResultSet executeQuery(String sql) throws SQLException { - execute(sql); - if (currentResultSet == null) { + StmtResult result = executeSql(sql); + if (result.cols() == null || result.cols().isEmpty()) { throw new SQLException("Query did not produce a ResultSet: " + sql); } + this.currentResultSet = new LibsqlResultSet(this, result, sql, false); + this.updateCount = -1; return currentResultSet; } @@ -48,10 +54,8 @@ public int executeUpdate(String sql) throws SQLException { @Override public boolean execute(String sql) throws SQLException { - checkClosed(); - closeCurrentResultSet(); - checkClosed(); - StmtResult result = connection.executeStatement(new Stmt(sql)); + StmtResult result = executeSql(sql); + updateGeneratedKeys(sql, result); if (result.cols() != null && !result.cols().isEmpty()) { this.currentResultSet = new LibsqlResultSet(this, result, sql, false); this.updateCount = -1; @@ -63,6 +67,13 @@ public boolean execute(String sql) throws SQLException { } } + private StmtResult executeSql(String sql) throws SQLException { + checkClosed(); + closeCurrentResultSet(); + checkClosed(); + return connection.executeStatement(new Stmt(sql)); + } + protected void closeCurrentResultSet() throws SQLException { if (currentResultSet != null) { currentResultSet.close(); @@ -70,6 +81,121 @@ protected void closeCurrentResultSet() throws SQLException { } } + protected void updateGeneratedKeys(String sql, StmtResult result) throws SQLException { + clearGeneratedKeys(); + if (sql != null && 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); + } + } + + protected void clearGeneratedKeys() throws SQLException { + if (generatedKeys != null) { + generatedKeys.close(); + generatedKeys = null; + } + } + + private static boolean isInsertSql(String sql) { + int start = skipWhitespace(sql, 0); + if (keywordAt(sql, start, "INSERT") || keywordAt(sql, start, "REPLACE")) { + return true; + } + if (!keywordAt(sql, start, "WITH")) { + return false; + } + int afterWith = start + "WITH".length(); + if (afterWith >= sql.length() || !Character.isWhitespace(sql.charAt(afterWith))) { + return false; + } + + int depth = 0; + int i = afterWith; + while (i < sql.length()) { + char c = sql.charAt(i); + if (c == '-' && i + 1 < sql.length() && sql.charAt(i + 1) == '-') { + i += 2; + while (i < sql.length() && sql.charAt(i) != '\n') { + i++; + } + } else if (c == '/' && i + 1 < sql.length() && sql.charAt(i + 1) == '*') { + i += 2; + while (i + 1 < sql.length() + && !(sql.charAt(i) == '*' && sql.charAt(i + 1) == '/')) { + i++; + } + i = Math.min(sql.length(), i + 2); + } else if (c == '\'' || c == '"' || c == '`') { + i = skipQuoted(sql, i, c); + } else if (c == '[') { + i++; + while (i < sql.length() && sql.charAt(i) != ']') { + i++; + } + if (i < sql.length()) { + i++; + } + } else if (c == '(') { + depth++; + i++; + } else if (c == ')') { + depth = Math.max(0, depth - 1); + i++; + if (depth == 0) { + int next = skipWhitespace(sql, i); + if (keywordAt(sql, next, "INSERT") || keywordAt(sql, next, "REPLACE")) { + return true; + } + } + } else { + i++; + } + } + return false; + } + + private static int skipWhitespace(String sql, int start) { + int i = start; + while (i < sql.length() && Character.isWhitespace(sql.charAt(i))) { + i++; + } + return i; + } + + private static int skipQuoted(String sql, int start, char quote) { + int i = start + 1; + while (i < sql.length()) { + if (sql.charAt(i) == quote) { + if (i + 1 < sql.length() && sql.charAt(i + 1) == quote) { + i += 2; + continue; + } + return i + 1; + } + i++; + } + return sql.length(); + } + + private static boolean keywordAt(String sql, int index, String keyword) { + return index + keyword.length() <= sql.length() + && sql.regionMatches(true, index, keyword, 0, keyword.length()); + } + + 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 { checkClosed(); @@ -128,6 +254,7 @@ public Connection getConnection() throws SQLException { public void close() throws SQLException { if (!closed) { closed = true; + clearGeneratedKeys(); if (currentResultSet != null) { currentResultSet.close(); currentResultSet = null; @@ -158,7 +285,14 @@ public boolean isClosed() throws SQLException { @Override public int getResultSetConcurrency() throws SQLException { return ResultSet.CONCUR_READ_ONLY; } @Override public int getResultSetType() throws SQLException { return ResultSet.TYPE_FORWARD_ONLY; } @Override public boolean getMoreResults(int current) throws SQLException { return getMoreResults(); } - @Override public ResultSet getGeneratedKeys() throws SQLException { throw new SQLException("getGeneratedKeys not supported directly"); } + @Override + public ResultSet getGeneratedKeys() throws SQLException { + checkClosed(); + if (generatedKeys == null) { + generatedKeys = emptyGeneratedKeys(); + } + return generatedKeys; + } @Override public int executeUpdate(String sql, int autoGeneratedKeys) throws SQLException { return executeUpdate(sql); } @Override public int executeUpdate(String sql, int[] columnIndexes) throws SQLException { return executeUpdate(sql); } @Override public int executeUpdate(String sql, String[] columnNames) throws SQLException { return executeUpdate(sql); } 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 new file mode 100644 index 0000000..adb3771 --- /dev/null +++ b/libsql-jdbc/src/test/java/com/github/magnusp/libsql/jdbc/LibsqlGeneratedKeysTest.java @@ -0,0 +1,206 @@ +package com.github.magnusp.libsql.jdbc; + +import com.github.magnusp.libsql.hrana.codec.Col; +import com.github.magnusp.libsql.hrana.codec.Row; +import com.github.magnusp.libsql.hrana.codec.Stmt; +import com.github.magnusp.libsql.hrana.codec.StmtResult; +import com.github.magnusp.libsql.hrana.codec.Value; +import org.junit.jupiter.api.Test; + +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayDeque; +import java.util.Deque; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +class LibsqlGeneratedKeysTest { + + private static StmtResult updateResult(long affectedRows, Long lastInsertRowid) { + return new StmtResult(List.of(), List.of(), affectedRows, lastInsertRowid); + } + + private static StmtResult insertResult(long rowid) { + return updateResult(1, rowid); + } + + @Test + void statementInsertReturnsGeneratedKey() throws SQLException { + RecordingConnection connection = new RecordingConnection(insertResult(42L)); + + try (Statement statement = connection.createStatement()) { + statement.executeUpdate("insert into k (v) values ('a')"); + + try (ResultSet keys = statement.getGeneratedKeys()) { + assertThat(keys.getMetaData().getColumnCount()).isEqualTo(1); + assertThat(keys.next()).isTrue(); + assertThat(keys.getLong(1)).isEqualTo(42L); + assertThat(keys.next()).isFalse(); + } + } + } + + @Test + void statementReplaceReturnsGeneratedKey() throws SQLException { + RecordingConnection connection = new RecordingConnection(insertResult(7L)); + + try (Statement statement = connection.createStatement()) { + statement.executeUpdate("replace into k (id, v) values (7, 'a')"); + + try (ResultSet keys = statement.getGeneratedKeys()) { + assertThat(keys.next()).isTrue(); + assertThat(keys.getLong(1)).isEqualTo(7L); + } + } + } + + @Test + void cteInsertReturnsGeneratedKey() throws SQLException { + RecordingConnection connection = new RecordingConnection(insertResult(13L)); + + try (Statement statement = connection.createStatement()) { + statement.executeUpdate( + "with cte as (select lower('INSERT') as v) insert into k (v) select v from cte"); + + try (ResultSet keys = statement.getGeneratedKeys()) { + assertThat(keys.next()).isTrue(); + assertThat(keys.getLong(1)).isEqualTo(13L); + } + } + } + + @Test + void cteSelectReturnsEmptyGeneratedKeys() throws SQLException { + RecordingConnection connection = new RecordingConnection(updateResult(1, 99L)); + + try (Statement statement = connection.createStatement()) { + statement.executeUpdate("with cte as (select 'insert' as v) select v from cte"); + + try (ResultSet keys = statement.getGeneratedKeys()) { + assertThat(keys.next()).isFalse(); + } + } + } + + @Test + void preparedStatementInsertReturnsGeneratedKey() throws SQLException { + RecordingConnection connection = new RecordingConnection(insertResult(5L)); + + try (LibsqlPreparedStatement statement = + new LibsqlPreparedStatement(connection, "insert into k (v) values (?)")) { + statement.setString(1, "a"); + statement.executeUpdate(); + + try (ResultSet keys = statement.getGeneratedKeys()) { + assertThat(keys.next()).isTrue(); + assertThat(keys.getLong(1)).isEqualTo(5L); + assertThat(keys.next()).isFalse(); + } + } + } + + @Test + void preparedStatementInsertReturningReturnsGeneratedKey() throws SQLException { + RecordingConnection connection = new RecordingConnection( + new StmtResult( + List.of(new Col("id", "INTEGER")), + List.of(new Row(List.of(Value.of(9L)))), + 1, + 9L)); + + try (LibsqlPreparedStatement statement = + new LibsqlPreparedStatement(connection, "insert into k (v) values (?) returning id")) { + statement.setString(1, "a"); + statement.execute(); + + try (ResultSet keys = statement.getGeneratedKeys()) { + assertThat(keys.next()).isTrue(); + assertThat(keys.getLong(1)).isEqualTo(9L); + } + } + } + + @Test + void getGeneratedKeysWithoutExecuteIsEmpty() throws SQLException { + RecordingConnection connection = new RecordingConnection(insertResult(1L)); + + try (Statement statement = connection.createStatement(); + ResultSet keys = statement.getGeneratedKeys()) { + assertThat(keys.getMetaData().getColumnCount()).isEqualTo(1); + assertThat(keys.next()).isFalse(); + } + } + + @Test + void nonInsertExecuteReturnsEmptyGeneratedKeys() throws SQLException { + RecordingConnection connection = new RecordingConnection(updateResult(1, 99L)); + + try (Statement statement = connection.createStatement()) { + statement.executeUpdate("update k set v = 'b'"); + + try (ResultSet keys = statement.getGeneratedKeys()) { + assertThat(keys.getMetaData().getColumnCount()).isEqualTo(1); + assertThat(keys.next()).isFalse(); + } + } + } + + @Test + void laterNonInsertClearsPreviousGeneratedKey() throws SQLException { + RecordingConnection connection = new RecordingConnection(insertResult(11L), updateResult(1, 99L)); + + try (Statement statement = connection.createStatement()) { + statement.executeUpdate("insert into k (v) values ('a')"); + statement.executeUpdate("delete from k"); + + try (ResultSet keys = statement.getGeneratedKeys()) { + assertThat(keys.next()).isFalse(); + } + } + } + + @Test + void executeQueryDoesNotClearPreviousGeneratedKey() throws SQLException { + StmtResult rows = new StmtResult( + List.of(new Col("id", "INTEGER")), + List.of(new Row(List.of(Value.of(3L)))), + 0, + null); + RecordingConnection connection = new RecordingConnection(insertResult(3L), rows); + + try (Statement statement = connection.createStatement()) { + statement.executeUpdate("insert into k (v) values ('a')"); + try (ResultSet resultSet = statement.executeQuery("select id from k")) { + assertThat(resultSet.next()).isTrue(); + } + + try (ResultSet keys = statement.getGeneratedKeys()) { + assertThat(keys.next()).isTrue(); + assertThat(keys.getLong(1)).isEqualTo(3L); + } + } + } + + @Test + void supportsGetGeneratedKeysMatchesImplementedBehavior() throws SQLException { + LibsqlDatabaseMetaData metadata = new LibsqlDatabaseMetaData(new RecordingConnection(insertResult(1L))); + + assertThat(metadata.supportsGetGeneratedKeys()).isTrue(); + } + + private static final class RecordingConnection extends LibsqlConnection { + private final Deque results = new ArrayDeque<>(); + + private RecordingConnection(StmtResult... results) { + super(null, null); + this.results.addAll(List.of(results)); + } + + @Override + public StmtResult executeStatement(Stmt stmt) { + return results.size() == 1 ? results.peek() : results.remove(); + } + } +}