From 7dec6c676c999c2b684997e9d04fb6497bf00f81 Mon Sep 17 00:00:00 2001
From: duanzhengqiang <duanzhengqiang@apache.org>
Date: Mon, 23 Dec 2024 17:02:28 +0800
Subject: [PATCH] Support MySQL update with statement parse

---
 .../mysql/visitor/statement/MySQLStatementVisitor.java    | 3 +++
 .../parser/statement/mysql/dml/MySQLUpdateStatement.java  | 8 ++++++++
 2 files changed, 11 insertions(+)

diff --git a/parser/sql/dialect/mysql/src/main/java/org/apache/shardingsphere/sql/parser/mysql/visitor/statement/MySQLStatementVisitor.java b/parser/sql/dialect/mysql/src/main/java/org/apache/shardingsphere/sql/parser/mysql/visitor/statement/MySQLStatementVisitor.java
index 28fb3b1022bb4..c7d56d971a903 100644
--- a/parser/sql/dialect/mysql/src/main/java/org/apache/shardingsphere/sql/parser/mysql/visitor/statement/MySQLStatementVisitor.java
+++ b/parser/sql/dialect/mysql/src/main/java/org/apache/shardingsphere/sql/parser/mysql/visitor/statement/MySQLStatementVisitor.java
@@ -1552,6 +1552,9 @@ public ASTNode visitUpdate(final UpdateContext ctx) {
         if (null != ctx.limitClause()) {
             result.setLimit((LimitSegment) visit(ctx.limitClause()));
         }
+        if (null != ctx.withClause()) {
+            result.setWithSegment((WithSegment) visit(ctx.withClause()));
+        }
         result.addParameterMarkerSegments(getParameterMarkerSegments());
         return result;
     }
diff --git a/parser/sql/statement/type/mysql/src/main/java/org/apache/shardingsphere/sql/parser/statement/mysql/dml/MySQLUpdateStatement.java b/parser/sql/statement/type/mysql/src/main/java/org/apache/shardingsphere/sql/parser/statement/mysql/dml/MySQLUpdateStatement.java
index 57aea84c4d28b..8677efa3dfa2c 100644
--- a/parser/sql/statement/type/mysql/src/main/java/org/apache/shardingsphere/sql/parser/statement/mysql/dml/MySQLUpdateStatement.java
+++ b/parser/sql/statement/type/mysql/src/main/java/org/apache/shardingsphere/sql/parser/statement/mysql/dml/MySQLUpdateStatement.java
@@ -20,6 +20,7 @@
 import lombok.Setter;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.order.OrderBySegment;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.pagination.limit.LimitSegment;
+import org.apache.shardingsphere.sql.parser.statement.core.segment.generic.WithSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.statement.dml.UpdateStatement;
 import org.apache.shardingsphere.sql.parser.statement.mysql.MySQLStatement;
 
@@ -31,10 +32,17 @@
 @Setter
 public final class MySQLUpdateStatement extends UpdateStatement implements MySQLStatement {
     
+    private WithSegment withSegment;
+    
     private OrderBySegment orderBy;
     
     private LimitSegment limit;
     
+    @Override
+    public Optional<WithSegment> getWithSegment() {
+        return Optional.ofNullable(withSegment);
+    }
+    
     @Override
     public Optional<OrderBySegment> getOrderBy() {
         return Optional.ofNullable(orderBy);