feat(jdbc): implement getGeneratedKeys for inserts

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

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

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

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<Stmt> 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,13 +67,135 @@ 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();
             currentResultSet = null;
         }
     }
 
+    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); }