From 2cf1b268ad51c8baa8d502823bdde8efa6466a09 Mon Sep 17 00:00:00 2001
From: Stephen Lombardo <sjlombardo@zetetic.net>
Date: Tue, 23 Jan 2024 10:42:46 -0500
Subject: [PATCH] fix(jdbc): support CTE and REPLACE syntax with
 Statement#getGeneratedKeys

---
 .../java/org/sqlite/core/CoreStatement.java   | 10 ++++++-
 src/test/java/org/sqlite/StatementTest.java   | 30 +++++++++++++++++--
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/src/main/java/org/sqlite/core/CoreStatement.java b/src/main/java/org/sqlite/core/CoreStatement.java
index b5834b0dd..a492f04ca 100644
--- a/src/main/java/org/sqlite/core/CoreStatement.java
+++ b/src/main/java/org/sqlite/core/CoreStatement.java
@@ -18,6 +18,7 @@
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.regex.Pattern;
 import org.sqlite.SQLiteConnection;
 import org.sqlite.SQLiteConnectionConfig;
 import org.sqlite.jdbc3.JDBC3Connection;
@@ -37,6 +38,13 @@ public abstract class CoreStatement implements Codes {
     private Statement generatedKeysStat = null;
     private ResultSet generatedKeysRs = null;
 
+    // pattern for matching insert statements of the general format starting with INSERT or REPLACE.
+    // CTEs used prior to the insert or replace keyword are also be permitted.
+    private final Pattern insertPattern =
+            Pattern.compile(
+                    "^(with\\s+.+\\(.+?\\))*\\s*(insert|replace)",
+                    Pattern.DOTALL | Pattern.CASE_INSENSITIVE);
+
     protected CoreStatement(SQLiteConnection c) {
         conn = c;
         rs = new JDBC4ResultSet(this);
@@ -174,7 +182,7 @@ protected void clearGeneratedKeys() throws SQLException {
      */
     public void updateGeneratedKeys() throws SQLException {
         clearGeneratedKeys();
-        if (sql != null && sql.toLowerCase().startsWith("insert")) {
+        if (sql != null && insertPattern.matcher(sql.trim().toLowerCase()).find()) {
             generatedKeysStat = conn.createStatement();
             generatedKeysRs = generatedKeysStat.executeQuery("SELECT last_insert_rowid();");
         }
diff --git a/src/test/java/org/sqlite/StatementTest.java b/src/test/java/org/sqlite/StatementTest.java
index ee08f4c90..ebb50ec9d 100644
--- a/src/test/java/org/sqlite/StatementTest.java
+++ b/src/test/java/org/sqlite/StatementTest.java
@@ -305,22 +305,48 @@ public void closeOnFalseNext() throws SQLException {
     public void getGeneratedKeys() throws SQLException {
         ResultSet rs;
         stat.executeUpdate("create table t1 (c1 integer primary key, v);");
-        stat.executeUpdate("insert into t1 (v) values ('red');");
 
+        // test standard insert operation
+        stat.executeUpdate("insert into t1 (v) values ('red');");
         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();
 
+        // test INSERT ith special replace keyword. This will trigger a primary key conflict on the
+        // first
+        // inserted row ('red') and replace the record with a value of 'yellow' with the same
+        // primary
+        // key. The value returned from getGeneratedKeys should be primary key of the replaced
+        // record
+        stat.executeUpdate("replace into t1 (c1, v) values (1, 'yellow');");
+        rs = stat.getGeneratedKeys();
+        assertThat(rs.next()).isTrue();
+        assertThat(rs.getInt(1)).isEqualTo(1);
+        rs.close();
+
+        // test INSERT with common table expression
+        stat.executeUpdate(
+                "with colors as (select 'green' as color)\n"
+                        + "insert into t1 (v) select color from colors;");
+        rs = stat.getGeneratedKeys();
+        assertThat(rs.next()).isTrue();
+        assertThat(rs.getInt(1)).isEqualTo(3);
+        rs.close();
+
+        stat.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();
+        stat.executeUpdate(
+                "with colors as (select 'insert' as color) update t1 set v = (select color from colors);");
         rs = stat2.getGeneratedKeys();
         assertThat(rs).isNotNull();
         assertThat(rs.next()).isFalse();