From 47bf6b046025a641574fd435ef7b3ddaa390af11 Mon Sep 17 00:00:00 2001
From: Zhichun Wu <zhicwu@gmail.com>
Date: Wed, 19 Jan 2022 20:12:54 +0800
Subject: [PATCH] Consider CTE when parsing insert statement

---
 .../src/main/javacc/ClickHouseSqlParser.jj    |  2 +-
 .../jdbc/ClickHousePreparedStatementTest.java | 26 +++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/clickhouse-jdbc/src/main/javacc/ClickHouseSqlParser.jj b/clickhouse-jdbc/src/main/javacc/ClickHouseSqlParser.jj
index 2550df6df..27a84c036 100644
--- a/clickhouse-jdbc/src/main/javacc/ClickHouseSqlParser.jj
+++ b/clickhouse-jdbc/src/main/javacc/ClickHouseSqlParser.jj
@@ -529,7 +529,7 @@ void dataClause(): {} {
                 <RPAREN> { token_source.removePosition(ClickHouseSqlStatement.KEYWORD_VALUES_END); }
             )*
             ( settingsPart() )?
-        | (LOOKAHEAD(2) (<SELECT> { token_source.addPosition(token); } columnExprList() <FROM> 
+        | (LOOKAHEAD(2) ((withClause())? <SELECT> { token_source.addPosition(token); } columnExprList() <FROM> 
             <INPUT> <LPAREN> <STRING_LITERAL> { token_source.input = ClickHouseSqlUtils.unescape(token.image); } <RPAREN>)?
             <FORMAT> <IDENTIFIER> { token_source.format = token.image; } )? (anyExprList())?
     } catch (ParseException e) {
diff --git a/clickhouse-jdbc/src/test/java/com/clickhouse/jdbc/ClickHousePreparedStatementTest.java b/clickhouse-jdbc/src/test/java/com/clickhouse/jdbc/ClickHousePreparedStatementTest.java
index 30334ec6b..708bb7327 100644
--- a/clickhouse-jdbc/src/test/java/com/clickhouse/jdbc/ClickHousePreparedStatementTest.java
+++ b/clickhouse-jdbc/src/test/java/com/clickhouse/jdbc/ClickHousePreparedStatementTest.java
@@ -513,4 +513,30 @@ public void testQueryWithNamedParameter() throws SQLException {
             Assert.assertFalse(rs.next());
         }
     }
+
+    @Test(groups = "integration")
+    public void testInsertWithAndSelect() throws Exception {
+        try (ClickHouseConnection conn = newConnection(new Properties());
+                Statement s = conn.createStatement()) {
+            s.execute("drop table if exists test_insert_with_and_select; "
+                    + "CREATE TABLE test_insert_with_and_select(value String) ENGINE=Memory");
+            try (PreparedStatement ps = conn.prepareStatement(
+                    "INSERT INTO test_insert_with_and_select(value) WITH t as ( SELECT 'testValue1') SELECT * FROM t")) {
+                ps.executeUpdate();
+            }
+
+            try (PreparedStatement ps = conn.prepareStatement(
+                    "INSERT INTO test_insert_with_and_select(value) WITH t as ( SELECT 'testValue2' as value) SELECT * FROM t WHERE value != ?")) {
+                ps.setString(1, "");
+                ps.executeUpdate();
+            }
+
+            ResultSet rs = s.executeQuery("select * from test_insert_with_and_select order by value");
+            Assert.assertTrue(rs.next());
+            Assert.assertEquals(rs.getString("Value"), "testValue1");
+            Assert.assertTrue(rs.next());
+            Assert.assertEquals(rs.getString("VALUE"), "testValue2");
+            Assert.assertFalse(rs.next());
+        }
+    }
 }