From 6c379c8b81e2e4fd6a8f497f2bfce5c79d3fa6a1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=AB=98=E9=93=81?= <shaojin.wensj@alibaba-inc.com>
Date: Mon, 7 Jun 2021 01:23:02 +0800
Subject: [PATCH] improved mysql parser, update support cte fix #4289

---
 .../mysql/visitor/MySqlOutputVisitor.java     |  5 +++
 .../druid/sql/parser/SQLStatementParser.java  |  4 +++
 .../sql/visitor/SQLASTOutputVisitor.java      |  5 +++
 .../sql/mysql/visitor/MySqlResourceTest.java  |  4 +++
 src/test/resources/bvt/parser/mysql-33.txt    | 31 +++++++++++++++++++
 5 files changed, 49 insertions(+)
 create mode 100644 src/test/resources/bvt/parser/mysql-33.txt

diff --git a/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlOutputVisitor.java b/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlOutputVisitor.java
index a209f0c922..03a03a255f 100644
--- a/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlOutputVisitor.java
+++ b/src/main/java/com/alibaba/druid/sql/dialect/mysql/visitor/MySqlOutputVisitor.java
@@ -2236,6 +2236,11 @@ public class MySqlOutputVisitor extends SQLASTOutputVisitor implements MySqlASTV
             }
         }
 
+        if (x.getWith() != null) {
+            x.getWith().accept(this);
+            println();
+        }
+
         List<SQLExpr> returning = x.getReturning();
         if (returning != null && returning.size() > 0) {
             print0(ucase ? "SELECT " : "select ");
diff --git a/src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java b/src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java
index f0d78258cd..197fd31bfc 100644
--- a/src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java
+++ b/src/main/java/com/alibaba/druid/sql/parser/SQLStatementParser.java
@@ -5257,6 +5257,10 @@ public class SQLStatementParser extends SQLParser {
             HiveMultiInsertStatement insert = (HiveMultiInsertStatement) this.parseInsert();
             insert.setWith(with);
             return insert;
+        } else if (lexer.token == UPDATE) {
+            SQLUpdateStatement update = this.parseUpdateStatement();
+            update.setWith(with);
+            return update;
         }
 
         throw new ParserException("TODO. " + lexer.info());
diff --git a/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java b/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java
index ec39a24486..699f556921 100644
--- a/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java
+++ b/src/main/java/com/alibaba/druid/sql/visitor/SQLASTOutputVisitor.java
@@ -3468,6 +3468,11 @@ public class SQLASTOutputVisitor extends SQLASTVisitorAdapter implements Paramet
             }
         }
 
+        if (x.getWith() != null) {
+            x.getWith().accept(this);
+            println();
+        }
+
         print0(ucase ? "UPDATE " : "update ");
 
         printTableSource(x.getTableSource());
diff --git a/src/test/java/com/alibaba/druid/bvt/sql/mysql/visitor/MySqlResourceTest.java b/src/test/java/com/alibaba/druid/bvt/sql/mysql/visitor/MySqlResourceTest.java
index fa975cd8d3..53a57c5752 100644
--- a/src/test/java/com/alibaba/druid/bvt/sql/mysql/visitor/MySqlResourceTest.java
+++ b/src/test/java/com/alibaba/druid/bvt/sql/mysql/visitor/MySqlResourceTest.java
@@ -123,6 +123,10 @@ public class MySqlResourceTest extends TestCase {
         exec_test("bvt/parser/mysql-32.txt");
     }
 
+    public void test_33() throws Exception {
+        exec_test("bvt/parser/mysql-33.txt");
+    }
+
     public void exec_test(String resource) throws Exception {
 //        System.out.println(resource);
         InputStream is = null;
diff --git a/src/test/resources/bvt/parser/mysql-33.txt b/src/test/resources/bvt/parser/mysql-33.txt
new file mode 100644
index 0000000000..fa8a530358
--- /dev/null
+++ b/src/test/resources/bvt/parser/mysql-33.txt
@@ -0,0 +1,31 @@
+with recursive tti as (
+    select t.task_id
+    from task_info t
+    where t.task_id in (?)
+    union
+    select t.task_id
+    from task_info t
+    inner join tti tt on t.pid = tt.task_id
+)
+update task_info t
+join (select now() as 'now') n
+inner join tti tt on tt.task_id = t.task_id
+set t.is_del = 1,
+t.delete_time = n.now,
+t.delete_clerk_id = ?
+---------------------------
+WITH RECURSIVE tti AS (
+		SELECT t.task_id
+		FROM task_info t
+		WHERE t.task_id IN (?)
+		UNION
+		SELECT t.task_id
+		FROM task_info t
+			INNER JOIN tti tt ON t.pid = tt.task_id
+	)
+UPDATE task_info t
+	JOIN (
+		SELECT now() AS "now"
+	) n
+	INNER JOIN tti tt ON tt.task_id = t.task_id
+SET t.is_del = 1, t.delete_time = n.now, t.delete_clerk_id = ?
\ No newline at end of file