From d463ef7a6552bf25754e86909fe1bc4060e53a3e Mon Sep 17 00:00:00 2001
From: Gauthier Roebroeck <gauthier.roebroeck@gmail.com>
Date: Thu, 26 Sep 2024 14:31:24 +0800
Subject: [PATCH] feat(jdbc): add setGetGeneratedKeys on SQLiteDataSource and
 SQLiteConfig for proper cascading

Refs: #1135
---
 USAGE.md                                      |  3 +-
 src/main/java/org/sqlite/SQLiteConfig.java    |  8 +++++
 .../java/org/sqlite/SQLiteDataSource.java     |  9 +++++
 .../java/org/sqlite/SQLiteConfigTest.java     |  3 ++
 .../java/org/sqlite/SQLiteDataSourceTest.java | 33 +++++++++++++++++++
 5 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/USAGE.md b/USAGE.md
index 4e879a626b..aa183d3e92 100644
--- a/USAGE.md
+++ b/USAGE.md
@@ -147,7 +147,8 @@ SQLite has limited support to retrieve generated keys, using [last_insert_rowid]
 
 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:
+You can disable the retrieval of generated keys in 3 ways:
+- via `SQLiteDataSource#setGetGeneratedKeys(false)`
 - via `SQLiteConnectionConfig#setGetGeneratedKeys(false)`:
 - using the pragma `jdbc.get_generated_keys`:
 ```java
diff --git a/src/main/java/org/sqlite/SQLiteConfig.java b/src/main/java/org/sqlite/SQLiteConfig.java
index e18babe942..0c4cf38434 100755
--- a/src/main/java/org/sqlite/SQLiteConfig.java
+++ b/src/main/java/org/sqlite/SQLiteConfig.java
@@ -1189,4 +1189,12 @@ public void setBusyTimeout(int milliseconds) {
     public int getBusyTimeout() {
         return busyTimeout;
     }
+
+    public boolean isGetGeneratedKeys() {
+        return this.defaultConnectionConfig.isGetGeneratedKeys();
+    }
+
+    public void setGetGeneratedKeys(boolean generatedKeys) {
+        this.defaultConnectionConfig.setGetGeneratedKeys(generatedKeys);
+    }
 }
diff --git a/src/main/java/org/sqlite/SQLiteDataSource.java b/src/main/java/org/sqlite/SQLiteDataSource.java
index 786b3ff8ee..a20c1da4fb 100755
--- a/src/main/java/org/sqlite/SQLiteDataSource.java
+++ b/src/main/java/org/sqlite/SQLiteDataSource.java
@@ -437,6 +437,15 @@ public void setTransactionMode(String transactionMode) {
         config.setTransactionMode(transactionMode);
     }
 
+    /**
+     * Configure where generated keys will be retrieved for this database.
+     *
+     * @param generatedKeys true to retrieve generated keys
+     */
+    public void setGetGeneratedKeys(boolean generatedKeys) {
+        config.setGetGeneratedKeys(generatedKeys);
+    }
+
     /**
      * Sets the value of the user-version. It is a big-endian 32-bit signed integer stored in the
      * database header at offset 60.
diff --git a/src/test/java/org/sqlite/SQLiteConfigTest.java b/src/test/java/org/sqlite/SQLiteConfigTest.java
index 24f2030ad0..b37b35a969 100644
--- a/src/test/java/org/sqlite/SQLiteConfigTest.java
+++ b/src/test/java/org/sqlite/SQLiteConfigTest.java
@@ -18,6 +18,7 @@ public void toProperties() {
         config.setDateStringFormat("yyyy/mm/dd");
         config.setDatePrecision("seconds");
         config.setDateClass("real");
+        config.setGetGeneratedKeys(false);
 
         Properties properties = config.toProperties();
 
@@ -27,6 +28,8 @@ public void toProperties() {
                 .isEqualTo(SQLiteConfig.DatePrecision.SECONDS.name());
         assertThat(properties.getProperty(SQLiteConfig.Pragma.DATE_CLASS.getPragmaName()))
                 .isEqualTo(SQLiteConfig.DateClass.REAL.name());
+        assertThat(properties.getProperty(Pragma.JDBC_GET_GENERATED_KEYS.getPragmaName()))
+                .isEqualTo("false");
     }
 
     @Test
diff --git a/src/test/java/org/sqlite/SQLiteDataSourceTest.java b/src/test/java/org/sqlite/SQLiteDataSourceTest.java
index b7674644cc..d01d630aeb 100755
--- a/src/test/java/org/sqlite/SQLiteDataSourceTest.java
+++ b/src/test/java/org/sqlite/SQLiteDataSourceTest.java
@@ -14,6 +14,7 @@
 import java.nio.ByteOrder;
 import java.sql.Connection;
 import java.sql.ResultSet;
+import java.sql.SQLException;
 import java.sql.Statement;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -92,4 +93,36 @@ public void setBusyTimeout() {
                 .isEqualTo("1234");
         assertThat(ds.getConfig().getBusyTimeout()).isEqualTo(1234);
     }
+
+    @Test
+    public void setGetGeneratedKeys() throws SQLException {
+        final SQLiteDataSource ds = new SQLiteDataSource();
+        ds.setGetGeneratedKeys(false);
+        assertThat(
+                        ds.getConfig()
+                                .toProperties()
+                                .getProperty(
+                                        SQLiteConfig.Pragma.JDBC_GET_GENERATED_KEYS.pragmaName))
+                .isEqualTo("false");
+        assertThat(ds.getConfig().isGetGeneratedKeys()).isEqualTo(false);
+        assertThat(
+                        ((SQLiteConnection) ds.getConnection())
+                                .getConnectionConfig()
+                                .isGetGeneratedKeys())
+                .isFalse();
+
+        ds.setGetGeneratedKeys(true);
+        assertThat(
+                        ds.getConfig()
+                                .toProperties()
+                                .getProperty(
+                                        SQLiteConfig.Pragma.JDBC_GET_GENERATED_KEYS.pragmaName))
+                .isEqualTo("true");
+        assertThat(ds.getConfig().isGetGeneratedKeys()).isEqualTo(true);
+        assertThat(
+                        ((SQLiteConnection) ds.getConnection())
+                                .getConnectionConfig()
+                                .isGetGeneratedKeys())
+                .isTrue();
+    }
 }