Support MySQL update with statement parse

apache · Dec 23, 2024 · 7dec6c6 · 7dec6c6
1 parent f246711
commit 7dec6c6
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/...a/org/apache/shardingsphere/sql/parser/mysql/visitor/statement/MySQLStatementVisitor.java b/...a/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/...n/java/org/apache/shardingsphere/sql/parser/statement/mysql/dml/MySQLUpdateStatement.java b/...n/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);