feat(jdbc): add configuration to disable eager retrieval of generated…

… keys Closes: xerial#1135
gotson · Sep 25, 2024 · 62b5ebb · 62b5ebb
1 parent 5db235f
commit 62b5ebb
Show file tree

Hide file tree

Showing 6 changed files with 95 additions and 12 deletions.
diff --git a/USAGE.md b/USAGE.md
@@ -139,6 +139,21 @@ You set the mode at the connection string level:
 try (Connection connection = DriverManager.getConnection("jdbc:sqlite:db.sqlite?hexkey_mode=sse", "", "AE...")) { /*...*/ }
 ```
 
+## Generated keys
+
+SQLite has limited support to retrieve generated keys, using [last_insert_rowid](https://www.sqlite.org/c3ref/last_insert_rowid.html), with the following limitations:
+- a single ID can be retrieved, even if multiple rows were added or updated
+- it needs to be called right after the statement
+
+By default the driver will eagerly retrieve the generated keys after each statement, which may impact performances.
+
+You can disable the retrieval of generated keys in 2 ways:
+- via `SQLiteConnectionConfig#setGetGeneratedKeys(false)`:
+- using the pragma `jdbc.get_generated_keys`:
+```java
+try (Connection connection = DriverManager.getConnection("jdbc:sqlite::memory:?jdbc.get_generated_keys=false")) { /*...*/ }
+```
+
 ## Explicit read only transactions (use with Hibernate)
 
 In order for the driver to be compliant with Hibernate, it needs to allow setting the read only flag after a connection has been created.

diff --git a/src/main/java/org/sqlite/SQLiteConfig.java b/src/main/java/org/sqlite/SQLiteConfig.java
@@ -188,6 +188,7 @@ public void apply(Connection conn) throws SQLException {
 
         // exclude this "fake" pragma from execution
         pragmaParams.remove(Pragma.JDBC_EXPLICIT_READONLY.pragmaName);
+        pragmaParams.remove(Pragma.JDBC_GET_GENERATED_KEYS.pragmaName);
 
         Statement stat = conn.createStatement();
         try {
@@ -330,6 +331,9 @@ public Properties toProperties() {
                 defaultConnectionConfig.getDateStringFormat());
         pragmaTable.setProperty(
                 Pragma.JDBC_EXPLICIT_READONLY.pragmaName, this.explicitReadOnly ? "true" : "false");
+        pragmaTable.setProperty(
+                Pragma.JDBC_GET_GENERATED_KEYS.pragmaName,
+                defaultConnectionConfig.isGetGeneratedKeys() ? "true" : "false");
         return pragmaTable;
     }
 
@@ -542,7 +546,9 @@ public enum Pragma {
 
         // extensions: "fake" pragmas to allow conformance with JDBC
         JDBC_EXPLICIT_READONLY(
-                "jdbc.explicit_readonly", "Set explicit read only transactions", null);
+                "jdbc.explicit_readonly", "Set explicit read only transactions", null),
+        JDBC_GET_GENERATED_KEYS(
+                "jdbc.get_generated_keys", "Enable retrieval of generated keys", OnOff.Values);
 
         public final String pragmaName;
         public final String[] choices;

diff --git a/src/main/java/org/sqlite/SQLiteConnectionConfig.java b/src/main/java/org/sqlite/SQLiteConnectionConfig.java
@@ -19,6 +19,7 @@ public class SQLiteConnectionConfig implements Cloneable {
     private int transactionIsolation = Connection.TRANSACTION_SERIALIZABLE;
     private SQLiteConfig.TransactionMode transactionMode = SQLiteConfig.TransactionMode.DEFERRED;
     private boolean autoCommit = true;
+    private boolean getGeneratedKeys = true;
 
     public static SQLiteConnectionConfig fromPragmaTable(Properties pragmaTable) {
         return new SQLiteConnectionConfig(
@@ -38,7 +39,10 @@ public static SQLiteConnectionConfig fromPragmaTable(Properties pragmaTable) {
                         pragmaTable.getProperty(
                                 SQLiteConfig.Pragma.TRANSACTION_MODE.pragmaName,
                                 SQLiteConfig.TransactionMode.DEFERRED.name())),
-                true);
+                true,
+                Boolean.parseBoolean(
+                        pragmaTable.getProperty(
+                                SQLiteConfig.Pragma.JDBC_GET_GENERATED_KEYS.pragmaName, "true")));
     }
 
     public SQLiteConnectionConfig(
@@ -47,13 +51,15 @@ public SQLiteConnectionConfig(
             String dateStringFormat,
             int transactionIsolation,
             SQLiteConfig.TransactionMode transactionMode,
-            boolean autoCommit) {
+            boolean autoCommit,
+            boolean getGeneratedKeys) {
         setDateClass(dateClass);
         setDatePrecision(datePrecision);
         setDateStringFormat(dateStringFormat);
         setTransactionIsolation(transactionIsolation);
         setTransactionMode(transactionMode);
         setAutoCommit(autoCommit);
+        setGetGeneratedKeys(getGeneratedKeys);
     }
 
     public SQLiteConnectionConfig copyConfig() {
@@ -63,7 +69,8 @@ public SQLiteConnectionConfig copyConfig() {
                 dateStringFormat,
                 transactionIsolation,
                 transactionMode,
-                autoCommit);
+                autoCommit,
+                getGeneratedKeys);
     }
 
     public long getDateMultiplier() {
@@ -124,8 +131,16 @@ public void setTransactionMode(SQLiteConfig.TransactionMode transactionMode) {
         this.transactionMode = transactionMode;
     }
 
+    public boolean isGetGeneratedKeys() {
+        return getGeneratedKeys;
+    }
+
+    public void setGetGeneratedKeys(boolean getGeneratedKeys) {
+        this.getGeneratedKeys = getGeneratedKeys;
+    }
+
     private static final Map<SQLiteConfig.TransactionMode, String> beginCommandMap =
-            new EnumMap<SQLiteConfig.TransactionMode, String>(SQLiteConfig.TransactionMode.class);
+            new EnumMap<>(SQLiteConfig.TransactionMode.class);
 
     static {
         beginCommandMap.put(SQLiteConfig.TransactionMode.DEFERRED, "begin;");

diff --git a/src/main/java/org/sqlite/core/CoreStatement.java b/src/main/java/org/sqlite/core/CoreStatement.java
@@ -181,10 +181,12 @@ protected void clearGeneratedKeys() throws SQLException {
      * block.
      */
     public void updateGeneratedKeys() throws SQLException {
-        clearGeneratedKeys();
-        if (sql != null && INSERT_PATTERN.matcher(sql).find()) {
-            generatedKeysStat = conn.createStatement();
-            generatedKeysRs = generatedKeysStat.executeQuery("SELECT last_insert_rowid();");
+        if (conn.getConnectionConfig().isGetGeneratedKeys()) {
+            clearGeneratedKeys();
+            if (sql != null && INSERT_PATTERN.matcher(sql).find()) {
+                generatedKeysStat = conn.createStatement();
+                generatedKeysRs = generatedKeysStat.executeQuery("SELECT last_insert_rowid();");
+            }
         }
     }
 

diff --git a/src/test/java/org/sqlite/PrepStmtTest.java b/src/test/java/org/sqlite/PrepStmtTest.java
@@ -81,6 +81,36 @@ public void update() throws SQLException {
         rs.close();
     }
 
+    @Test
+    public void pragmaGetGeneratedKeys() throws SQLException {
+        SQLiteConnection connection =
+                (SQLiteConnection)
+                        DriverManager.getConnection(
+                                "jdbc:sqlite::memory:?jdbc.get_generated_keys=false");
+        assertThat(connection.getConnectionConfig().isGetGeneratedKeys()).isFalse();
+    }
+
+    @Test
+    public void updateWithoutGeneratedKeys() throws SQLException {
+        Connection conn =
+                DriverManager.getConnection("jdbc:sqlite::memory:?jdbc.get_generated_keys=false");
+
+        assertThat(conn.prepareStatement("create table s1 (c1);").executeUpdate()).isEqualTo(0);
+        PreparedStatement prep = conn.prepareStatement("insert into s1 values (?);");
+        prep.setInt(1, 3);
+        assertThat(prep.executeUpdate()).isEqualTo(1);
+        assertThat(prep.getResultSet()).isNull();
+        prep.setInt(1, 5);
+        assertThat(prep.executeUpdate()).isEqualTo(1);
+        prep.setInt(1, 7);
+        assertThat(prep.executeUpdate()).isEqualTo(1);
+
+        ResultSet rsgk = prep.getGeneratedKeys();
+        assertThat(rsgk.next()).isFalse();
+        rsgk.close();
+        prep.close();
+    }
+
     @Test
     public void multiUpdate() throws SQLException {
         stat.executeUpdate("create table test (c1);");

diff --git a/src/test/java/org/sqlite/StatementTest.java b/src/test/java/org/sqlite/StatementTest.java
@@ -1,8 +1,6 @@
 package org.sqlite;
 
-import static org.assertj.core.api.Assertions.assertThat;
-import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
-import static org.assertj.core.api.Assertions.assertThatNoException;
+import static org.assertj.core.api.Assertions.*;
 import static org.assertj.core.data.Offset.offset;
 
 import java.lang.reflect.Method;
@@ -351,6 +349,23 @@ public void getGeneratedKeys() throws SQLException {
         assertThat(rs).isNotNull();
         assertThat(rs.next()).isFalse();
         stat2.close();
+
+        // disable then re-enable generated keys retrieval
+        try {
+            ((SQLiteConnection) conn).getConnectionConfig().setGetGeneratedKeys(false);
+            stat.executeUpdate("insert into t1 (v) values ('red');");
+            rs = stat.getGeneratedKeys();
+            assertThat(rs.next()).isFalse();
+
+            ((SQLiteConnection) conn).getConnectionConfig().setGetGeneratedKeys(true);
+
+            stat.executeUpdate("insert into t1 (v) values ('red');");
+            rs = stat.getGeneratedKeys();
+            assertThat(rs.next()).isTrue();
+            rs.close();
+        } finally {
+            ((SQLiteConnection) conn).getConnectionConfig().setGetGeneratedKeys(true);
+        }
     }
 
     @Test