Skip to content

Commit

Permalink
feat(jdbc): reintroduce improved support for Statement#getGeneratedKeys
Browse files Browse the repository at this point in the history
  • Loading branch information
sjlombardo committed Jan 19, 2024
1 parent 4e3520c commit f7d49f6
Show file tree
Hide file tree
Showing 7 changed files with 144 additions and 33 deletions.
48 changes: 48 additions & 0 deletions src/main/java/org/sqlite/core/CoreStatement.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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);
Expand Down Expand Up @@ -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")) {

This comment has been minimized.

Copy link
@gwenn

gwenn Jan 20, 2024

Does not work with CTE:
https://sqlite.org/lang_insert.html

This comment has been minimized.

Copy link
@gotson

gotson Jan 22, 2024

Collaborator

please don't comment on commits

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;
}
}
2 changes: 1 addition & 1 deletion src/main/java/org/sqlite/jdbc3/JDBC3DatabaseMetaData.java
Original file line number Diff line number Diff line change
Expand Up @@ -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() */
Expand Down
21 changes: 16 additions & 5 deletions src/main/java/org/sqlite/jdbc3/JDBC3PreparedStatement.java
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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() */
Expand Down
41 changes: 18 additions & 23 deletions src/main/java/org/sqlite/jdbc3/JDBC3Statement.java
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ protected JDBC3Statement(SQLiteConnection conn) {

/** @see java.sql.Statement#close() */
public void close() throws SQLException {
clearGeneratedKeys();
internalClose();
}

Expand All @@ -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;
}
});
}

Expand Down Expand Up @@ -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();
}
Expand Down Expand Up @@ -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().
*
Expand Down
1 change: 0 additions & 1 deletion src/test/java/org/sqlite/DBMetaDataTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
7 changes: 6 additions & 1 deletion src/test/java/org/sqlite/PrepStmtTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
57 changes: 55 additions & 2 deletions src/test/java/org/sqlite/StatementTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit f7d49f6

Please sign in to comment.