fix(jdbc): support CTE and REPLACE syntax with Statement#getGenerated…

…Keys

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();");
         }

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();