feat(jdbc): reintroduce improved support for Statement#getGeneratedKeys

Author: sjlombardo

src/main/java/org/sqlite/core/CoreStatement.java

@@ -17,6 +17,7 @@
 
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.Statement;
 import org.sqlite.SQLiteConnection;
 import org.sqlite.SQLiteConnectionConfig;
 import org.sqlite.jdbc3.JDBC3Connection;
@@ -33,6 +34,9 @@ public abstract class CoreStatement implements Codes {
     protected Object[] batch = null;
     protected boolean resultsWaiting = false;
 
+    private Statement generatedKeysStat = null;
+    private ResultSet generatedKeysRs = null;
+
     protected CoreStatement(SQLiteConnection c) {
         conn = c;
         rs = new JDBC4ResultSet(this);
@@ -149,4 +153,48 @@ protected void checkIndex(int index) throws SQLException {
             throw new SQLException("Parameter index is invalid");
         }
     }
+
+    protected void clearGeneratedKeys() throws SQLException {
+        if (generatedKeysRs != null && !generatedKeysRs.isClosed()) {
+            generatedKeysRs.close();
+        }
+        generatedKeysRs = null;
+        if (generatedKeysStat != null && !generatedKeysStat.isClosed()) {
+            generatedKeysStat.close();
+        }
+        generatedKeysStat = null;
+    }
+
+    /**
+     * SQLite's last_insert_rowid() function is DB-specific. However, in this implementation we
+     * ensure the Generated Key result set is statement-specific by executing the query immediately
+     * after an insert operation is performed. The caller is simply responsible for calling
+     * updateGeneratedKeys on the statement object right after execute in a synchronized(connection)
+     * block.
+     */
+    public void updateGeneratedKeys() throws SQLException {
+        clearGeneratedKeys();
+        if (sql != null && sql.toLowerCase().startsWith("insert")) {
+            generatedKeysStat = conn.createStatement();
+            generatedKeysRs = generatedKeysStat.executeQuery("SELECT last_insert_rowid();");
+        }
+    }
+
+    /**
+     * This implementation uses SQLite's last_insert_rowid function to obtain the row ID. It cannot
+     * provide multiple values when inserting multiple rows. Suggestion is to use a <a
+     * href=https://www.sqlite.org/lang_returning.html>RETURNING</a> clause instead.
+     *
+     * @see java.sql.Statement#getGeneratedKeys()
+     */
+    public ResultSet getGeneratedKeys() throws SQLException {
+        // getGeneratedKeys is required to return an EmptyResult set if the statement
+        // did not generate any keys. Thus, if the generateKeysResultSet is NULL, spin
+        // up a new result set without any contents by issuing a query with a false where condition
+        if (generatedKeysRs == null) {
+            generatedKeysStat = conn.createStatement();
+            generatedKeysRs = generatedKeysStat.executeQuery("SELECT 1 WHERE 1 = 2;");
+        }
+        return generatedKeysRs;
+    }
 }

src/main/java/org/sqlite/jdbc3/JDBC3DatabaseMetaData.java

@@ -559,7 +559,7 @@ public boolean supportsFullOuterJoins() throws SQLException {
 
     /** @see java.sql.DatabaseMetaData#supportsGetGeneratedKeys() */
     public boolean supportsGetGeneratedKeys() {
-        return false;
+        return true;
     }
 
     /** @see java.sql.DatabaseMetaData#supportsGroupBy() */

src/main/java/org/sqlite/jdbc3/JDBC3PreparedStatement.java

@@ -54,10 +54,13 @@ public boolean execute() throws SQLException {
                 () -> {
                     boolean success = false;
                     try {
-                        resultsWaiting =
-                                conn.getDatabase().execute(JDBC3PreparedStatement.this, batch);
-                        success = true;
-                        updateCount = getDatabase().changes();
+                        synchronized (conn) {
+                            resultsWaiting =
+                                    conn.getDatabase().execute(JDBC3PreparedStatement.this, batch);
+                            updateGeneratedKeys();
+                            success = true;
+                            updateCount = getDatabase().changes();
+                        }
                         return 0 != columnCount;
                     } finally {
                         if (!success && !pointer.isClosed()) pointer.safeRunConsume(DB::reset);
@@ -119,7 +122,15 @@ public long executeLargeUpdate() throws SQLException {
         }
 
         return this.withConnectionTimeout(
-                () -> conn.getDatabase().executeUpdate(JDBC3PreparedStatement.this, batch));
+                () -> {
+                    synchronized (conn) {
+                        long rc =
+                                conn.getDatabase()
+                                        .executeUpdate(JDBC3PreparedStatement.this, batch);
+                        updateGeneratedKeys();
+                        return rc;
+                    }
+                });
     }
 
     /** @see java.sql.PreparedStatement#addBatch() */

src/main/java/org/sqlite/jdbc3/JDBC3Statement.java

@@ -32,6 +32,7 @@ protected JDBC3Statement(SQLiteConnection conn) {
 
     /** @see java.sql.Statement#close() */
     public void close() throws SQLException {
+        clearGeneratedKeys();
         internalClose();
     }
 
@@ -49,12 +50,14 @@ public boolean execute(final String sql) throws SQLException {
                     }
 
                     JDBC3Statement.this.sql = sql;
-
-                    conn.getDatabase().prepare(JDBC3Statement.this);
-                    boolean result = exec();
-                    updateCount = getDatabase().changes();
-                    exhaustedResults = false;
-                    return result;
+                    synchronized (conn) {
+                        conn.getDatabase().prepare(JDBC3Statement.this);
+                        boolean result = exec();
+                        updateGeneratedKeys();
+                        updateCount = getDatabase().changes();
+                        exhaustedResults = false;
+                        return result;
+                    }
                 });
     }
 
@@ -126,13 +129,16 @@ public long executeLargeUpdate(String sql) throws SQLException {
                         ext.execute(db);
                     } else {
                         try {
-                            changes = db.total_changes();
-
-                            // directly invokes the exec API to support multiple SQL statements
-                            int statusCode = db._exec(sql);
-                            if (statusCode != SQLITE_OK) throw DB.newSQLException(statusCode, "");
+                            synchronized (db) {
+                                changes = db.total_changes();
+                                // directly invokes the exec API to support multiple SQL statements
+                                int statusCode = db._exec(sql);
+                                if (statusCode != SQLITE_OK)
+                                    throw DB.newSQLException(statusCode, "");
+                                updateGeneratedKeys();
+                                changes = db.total_changes() - changes;
+                            }
 
-                            changes = db.total_changes() - changes;
                         } finally {
                             internalClose();
                         }
@@ -350,17 +356,6 @@ public void setFetchDirection(int direction) throws SQLException {
         }
     }
 
-    /**
-     * SQLite's last_insert_rowid() function is DB-specific, not statement specific, and cannot
-     * provide multiple values when inserting multiple rows. Suggestion is to use a <a
-     * href=https://www.sqlite.org/lang_returning.html>RETURNING</a> clause instead.
-     *
-     * @see java.sql.Statement#getGeneratedKeys()
-     */
-    public ResultSet getGeneratedKeys() throws SQLException {
-        throw unsupported();
-    }
-
     /**
      * SQLite does not support multiple results from execute().
      *

src/test/java/org/sqlite/DBMetaDataTest.java

@@ -51,7 +51,6 @@ public void close() throws SQLException {
     public void getTables() throws SQLException {
         ResultSet rs = meta.getTables(null, null, null, null);
         assertThat(rs).isNotNull();
-
         stat.close();
 
         assertThat(rs.next()).isTrue();

src/test/java/org/sqlite/PrepStmtTest.java

@@ -67,8 +67,13 @@ public void update() throws SQLException {
         assertThat(prep.executeUpdate()).isEqualTo(1);
         prep.setInt(1, 7);
         assertThat(prep.executeUpdate()).isEqualTo(1);
-        prep.close();
 
+        ResultSet rsgk = prep.getGeneratedKeys();
+        assertThat(rsgk.next()).isTrue();
+        assertThat(rsgk.getInt(1)).isEqualTo(3);
+        rsgk.close();
+
+        prep.close();
         // check results with normal statement
         ResultSet rs = stat.executeQuery("select sum(c1) from s1;");
         assertThat(rs.next()).isTrue();

src/test/java/org/sqlite/StatementTest.java

@@ -307,8 +307,61 @@ public void getGeneratedKeys() throws SQLException {
         stat.executeUpdate("create table t1 (c1 integer primary key, v);");
         stat.executeUpdate("insert into t1 (v) values ('red');");
 
-        assertThatExceptionOfType(SQLFeatureNotSupportedException.class)
-                .isThrownBy(() -> stat.getGeneratedKeys());
+        rs = stat.getGeneratedKeys();
+        assertThat(rs.next()).isTrue();
+        assertThat(rs.getInt(1)).isEqualTo(1);
+        rs.close();
+        stat.executeUpdate("insert into t1 (v) values ('blue');");
+        rs = stat.getGeneratedKeys();
+        assertThat(rs.next()).isTrue();
+        assertThat(rs.getInt(1)).isEqualTo(2);
+        rs.close();
+
+        // generated keys are now attached to the statement. calling getGeneratedKeys
+        // on a statement that has not generated any should return an empty result set
+        stat.close();
+        Statement stat2 = conn.createStatement();
+        rs = stat2.getGeneratedKeys();
+        assertThat(rs).isNotNull();
+        assertThat(rs.next()).isFalse();
+        stat2.close();
+    }
+
+    @Test
+    public void getGeneratedKeysIsStatementSpecific() throws SQLException {
+        /* this test ensures that the results of getGeneratedKeys are tied to
+          a specific statement. To verify this, we create two separate Statement
+          objects and then execute inserts on both. We then make getGeneratedKeys()
+          calls and verify that the two separate expected values are returned.
+
+          Note that the old implementation of getGeneratedKeys was called lazily, so
+          the result of both getGeneratedKeys calls would be the same value, the row ID
+          of the last insert on the connection. As a result it was unsafe to use
+          with multiple statements or in a multithreaded application.
+        */
+        stat.executeUpdate("create table t1 (c1 integer primary key, v);");
+
+        ResultSet rs1;
+        Statement stat1 = conn.createStatement();
+        ResultSet rs2;
+        Statement stat2 = conn.createStatement();
+
+        stat1.executeUpdate("insert into t1 (v) values ('red');");
+        stat2.executeUpdate("insert into t1 (v) values ('blue');");
+
+        rs2 = stat2.getGeneratedKeys();
+        rs1 = stat1.getGeneratedKeys();
+
+        assertThat(rs1.next()).isTrue();
+        assertThat(rs1.getInt(1)).isEqualTo(1);
+        rs1.close();
+
+        assertThat(rs2.next()).isTrue();
+        assertThat(rs2.getInt(1)).isEqualTo(2);
+        rs2.close();
+
+        stat1.close();
+        stat2.close();
     }
 
     @Test