Support parsing SQL Server UPDATE dbo.Cities sql apache#29185

yydeng626 · Jan 31, 2024 · 2d84039 · 2d84039
1 parent df3882b
commit 2d84039
Show file tree

Hide file tree

Showing 13 changed files with 434 additions and 11 deletions.
diff --git a/parser/sql/dialect/sqlserver/src/main/antlr4/imports/sqlserver/BaseRule.g4 b/parser/sql/dialect/sqlserver/src/main/antlr4/imports/sqlserver/BaseRule.g4
@@ -169,7 +169,7 @@ serviceName
     ;
 
 columnName
-    : (owner DOT_)? (name | scriptVariableName)
+    : ((databaseName DOT_)? (owner DOT_))? (name | scriptVariableName)
     ;
 
 scriptVariableName

diff --git a/parser/sql/dialect/sqlserver/src/main/antlr4/imports/sqlserver/DMLStatement.g4 b/parser/sql/dialect/sqlserver/src/main/antlr4/imports/sqlserver/DMLStatement.g4
@@ -48,11 +48,11 @@ exec
     ;
 
 update
-    : withClause? UPDATE top? tableReferences setAssignmentsClause whereClause? (OPTION queryHint)?
+    : withClause? UPDATE top? tableReferences withTableHint? setAssignmentsClause whereClause? optionHint?
     ;
 
 assignment
-    : columnName EQ_ assignmentValue
+    : columnName (EQ_ | DOT_) assignmentValue
     ;
 
 setAssignmentsClause
@@ -69,7 +69,11 @@ assignmentValue
     ;
 
 delete
-    : withClause? DELETE top? (singleTableClause | multipleTablesClause) outputClause? whereClause? (OPTION queryHint)?
+    : withClause? DELETE top? (singleTableClause | multipleTablesClause) outputClause? whereClause? optionHint?
+    ;
+
+optionHint
+    : OPTION queryHint
     ;
 
 singleTableClause
@@ -222,7 +226,7 @@ queryHint
     | MAXDOP INT_NUM_
     | MAXRECURSION INT_NUM_
     | NO_PERFORMANCE_SPOOL
-    | OPTIMIZE FOR LP_ AT_ name (UNKNOWN | EQ_ identifier)* RP_
+    | LP_ OPTIMIZE FOR LP_ variableName (UNKNOWN | EQ_ literals)* RP_ RP_
     | OPTIMIZE FOR UNKNOWN
     | PARAMETERIZATION (SIMPLE | FORCED)
     | QUERYTRACEON INT_NUM_

diff --git a/...server/src/main/antlr4/org/apache/shardingsphere/sql/parser/autogen/SQLServerStatement.g4 b/...server/src/main/antlr4/org/apache/shardingsphere/sql/parser/autogen/SQLServerStatement.g4
@@ -19,6 +19,7 @@ grammar SQLServerStatement;
 
 import Comments, TCLStatement, StoreProcedure, DALStatement;
 
+
 execute
     : (select
     | insert

diff --git a/...ache/shardingsphere/sql/parser/sqlserver/visitor/statement/SQLServerStatementVisitor.java b/...ache/shardingsphere/sql/parser/sqlserver/visitor/statement/SQLServerStatementVisitor.java
@@ -82,6 +82,7 @@
 import org.apache.shardingsphere.sql.parser.autogen.SQLServerStatementParser.NumberLiteralsContext;
 import org.apache.shardingsphere.sql.parser.autogen.SQLServerStatementParser.OpenJsonFunctionContext;
 import org.apache.shardingsphere.sql.parser.autogen.SQLServerStatementParser.OpenRowSetFunctionContext;
+import org.apache.shardingsphere.sql.parser.autogen.SQLServerStatementParser.OptionHintContext;
 import org.apache.shardingsphere.sql.parser.autogen.SQLServerStatementParser.OrderByClauseContext;
 import org.apache.shardingsphere.sql.parser.autogen.SQLServerStatementParser.OrderByItemContext;
 import org.apache.shardingsphere.sql.parser.autogen.SQLServerStatementParser.OutputClauseContext;
@@ -205,6 +206,7 @@
 import org.apache.shardingsphere.sql.parser.sql.common.value.literal.impl.StringLiteralValue;
 import org.apache.shardingsphere.sql.parser.sql.common.value.parametermarker.ParameterMarkerValue;
 import org.apache.shardingsphere.sql.parser.sql.dialect.segment.sqlserver.exec.ExecSegment;
+import org.apache.shardingsphere.sql.parser.sql.dialect.segment.sqlserver.hint.OptionHintSegment;
 import org.apache.shardingsphere.sql.parser.sql.dialect.segment.sqlserver.hint.TableHintLimitedSegment;
 import org.apache.shardingsphere.sql.parser.sql.dialect.segment.sqlserver.hint.WithTableHintSegment;
 import org.apache.shardingsphere.sql.parser.sql.dialect.segment.sqlserver.json.JsonNullClauseSegment;
@@ -353,7 +355,12 @@ public final ASTNode visitColumnName(final ColumnNameContext ctx) {
         }
         OwnerContext owner = ctx.owner();
         if (null != owner) {
-            result.setOwner(new OwnerSegment(owner.getStart().getStartIndex(), owner.getStop().getStopIndex(), (IdentifierValue) visit(owner.identifier())));
+            OwnerSegment ownerSegment = new OwnerSegment(owner.getStart().getStartIndex(), owner.getStop().getStopIndex(), (IdentifierValue) visit(owner.identifier()));
+            if (null != ctx.databaseName()) {
+                ownerSegment.setOwner(new OwnerSegment(ctx.databaseName().getStart().getStartIndex(), ctx.databaseName().getStop().getStopIndex(),
+                        (IdentifierValue) visit(ctx.databaseName().identifier())));
+            }
+            result.setOwner(ownerSegment);
         }
         return result;
     }
@@ -1248,20 +1255,35 @@ public ASTNode visitUpdate(final UpdateContext ctx) {
         }
         result.setTable((TableSegment) visit(ctx.tableReferences()));
         result.setSetAssignment((SetAssignmentSegment) visit(ctx.setAssignmentsClause()));
+        if (null != ctx.withTableHint()) {
+            result.setWithTableHintSegment((WithTableHintSegment) visit(ctx.withTableHint()));
+        }
         if (null != ctx.whereClause()) {
             result.setWhere((WhereSegment) visit(ctx.whereClause()));
         }
+        if (null != ctx.optionHint()) {
+            result.setOptionHintSegment((OptionHintSegment) visit(ctx.optionHint()));
+        }
         result.addParameterMarkerSegments(getParameterMarkerSegments());
         return result;
     }
 
+    @Override
+    public ASTNode visitOptionHint(final OptionHintContext ctx) {
+        return new OptionHintSegment(ctx.getStart().getStartIndex(), ctx.getStop().getStopIndex(), getOriginalText(ctx));
+    }
+
     @Override
     public ASTNode visitSetAssignmentsClause(final SetAssignmentsClauseContext ctx) {
         Collection<ColumnAssignmentSegment> assignments = new LinkedList<>();
         for (AssignmentContext each : ctx.assignment()) {
             assignments.add((ColumnAssignmentSegment) visit(each));
         }
-        return new SetAssignmentSegment(ctx.getStart().getStartIndex(), ctx.getStop().getStopIndex(), assignments);
+        SetAssignmentSegment result = new SetAssignmentSegment(ctx.getStart().getStartIndex(), ctx.getStop().getStopIndex(), assignments);
+        if (null != ctx.fromClause()) {
+            result.setFrom((TableSegment) visit(ctx.fromClause()));
+        }
+        return result;
     }
 
     @Override
@@ -1279,9 +1301,7 @@ public ASTNode visitAssignment(final AssignmentContext ctx) {
         List<ColumnSegment> columnSegments = new LinkedList<>();
         columnSegments.add(column);
         ExpressionSegment value = (ExpressionSegment) visit(ctx.assignmentValue());
-        ColumnAssignmentSegment result = new ColumnAssignmentSegment(ctx.getStart().getStartIndex(), ctx.getStop().getStopIndex(), columnSegments, value);
-        result.getColumns().add(column);
-        return result;
+        return new ColumnAssignmentSegment(ctx.getStart().getStartIndex(), ctx.getStop().getStopIndex(), columnSegments, value);
     }
 
     @Override

diff --git a/...che/shardingsphere/sql/parser/sql/common/segment/dml/assignment/SetAssignmentSegment.java b/...che/shardingsphere/sql/parser/sql/common/segment/dml/assignment/SetAssignmentSegment.java
@@ -19,7 +19,9 @@
 
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
+import lombok.Setter;
 import org.apache.shardingsphere.sql.parser.sql.common.segment.SQLSegment;
+import org.apache.shardingsphere.sql.parser.sql.common.segment.generic.table.TableSegment;
 
 import java.util.Collection;
 
@@ -35,4 +37,7 @@ public final class SetAssignmentSegment implements SQLSegment {
     private final int stopIndex;
 
     private final Collection<ColumnAssignmentSegment> assignments;
+
+    @Setter
+    private TableSegment from;
 }
diff --git a/.../org/apache/shardingsphere/sql/parser/sql/dialect/handler/dml/UpdateStatementHandler.java b/.../org/apache/shardingsphere/sql/parser/sql/dialect/handler/dml/UpdateStatementHandler.java
@@ -25,6 +25,7 @@
 import org.apache.shardingsphere.sql.parser.sql.common.segment.generic.WithSegment;
 import org.apache.shardingsphere.sql.parser.sql.common.statement.dml.UpdateStatement;
 import org.apache.shardingsphere.sql.parser.sql.dialect.handler.SQLStatementHandler;
+import org.apache.shardingsphere.sql.parser.sql.dialect.segment.sqlserver.hint.OptionHintSegment;
 import org.apache.shardingsphere.sql.parser.sql.dialect.statement.mysql.MySQLStatement;
 import org.apache.shardingsphere.sql.parser.sql.dialect.statement.mysql.dml.MySQLUpdateStatement;
 import org.apache.shardingsphere.sql.parser.sql.dialect.statement.oracle.dml.OracleUpdateStatement;
@@ -91,6 +92,19 @@ public static Optional<WhereSegment> getDeleteWhereSegment(final UpdateStatement
         return Optional.empty();
     }
 
+    /**
+     * Get option hint segment.
+     *
+     * @param updateStatement update statement
+     * @return option hint segment
+     */
+    public static Optional<OptionHintSegment> getOptionHintSegment(final UpdateStatement updateStatement) {
+        if (updateStatement instanceof SQLServerStatement) {
+            return ((SQLServerUpdateStatement) updateStatement).getOptionHintSegment();
+        }
+        return Optional.empty();
+    }
+
     /**
      * Set order by segment.
      * 

diff --git a/...pache/shardingsphere/sql/parser/sql/dialect/segment/sqlserver/hint/OptionHintSegment.java b/...pache/shardingsphere/sql/parser/sql/dialect/segment/sqlserver/hint/OptionHintSegment.java
@@ -0,0 +1,36 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.shardingsphere.sql.parser.sql.dialect.segment.sqlserver.hint;
+
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import org.apache.shardingsphere.sql.parser.sql.common.segment.SQLSegment;
+
+/**
+ * Option hint segment.
+ */
+@RequiredArgsConstructor
+@Getter
+public final class OptionHintSegment implements SQLSegment {
+
+    private final int startIndex;
+
+    private final int stopIndex;
+
+    private final String text;
+}
diff --git a/...ardingsphere/sql/parser/sql/dialect/statement/sqlserver/dml/SQLServerUpdateStatement.java b/...ardingsphere/sql/parser/sql/dialect/statement/sqlserver/dml/SQLServerUpdateStatement.java
@@ -20,6 +20,8 @@
 import lombok.Setter;
 import org.apache.shardingsphere.sql.parser.sql.common.segment.generic.WithSegment;
 import org.apache.shardingsphere.sql.parser.sql.common.statement.dml.UpdateStatement;
+import org.apache.shardingsphere.sql.parser.sql.dialect.segment.sqlserver.hint.OptionHintSegment;
+import org.apache.shardingsphere.sql.parser.sql.dialect.segment.sqlserver.hint.WithTableHintSegment;
 import org.apache.shardingsphere.sql.parser.sql.dialect.statement.sqlserver.SQLServerStatement;
 
 import java.util.Optional;
@@ -32,6 +34,10 @@ public final class SQLServerUpdateStatement extends UpdateStatement implements S
 
     private WithSegment withSegment;
 
+    private WithTableHintSegment withTableHintSegment;
+
+    private OptionHintSegment optionHintSegment;
+
     /**
      * Get with segment.
      *
@@ -40,4 +46,22 @@ public final class SQLServerUpdateStatement extends UpdateStatement implements S
     public Optional<WithSegment> getWithSegment() {
         return Optional.ofNullable(withSegment);
     }
+
+    /**
+     * Get with table hint segment.
+     *
+     * @return with table hint segment.
+     */
+    public Optional<WithTableHintSegment> getWithTableHintSegment() {
+        return Optional.ofNullable(withTableHintSegment);
+    }
+
+    /**
+     * Get option hint segment.
+     *
+     * @return option hint segment.
+     */
+    public Optional<OptionHintSegment> getOptionHintSegment() {
+        return Optional.ofNullable(optionHintSegment);
+    }
 }
diff --git a/...gsphere/test/it/sql/parser/internal/asserts/statement/dml/impl/UpdateStatementAssert.java b/...gsphere/test/it/sql/parser/internal/asserts/statement/dml/impl/UpdateStatementAssert.java
@@ -23,6 +23,7 @@
 import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.pagination.limit.LimitSegment;
 import org.apache.shardingsphere.sql.parser.sql.common.statement.dml.UpdateStatement;
 import org.apache.shardingsphere.sql.parser.sql.dialect.handler.dml.UpdateStatementHandler;
+import org.apache.shardingsphere.sql.parser.sql.dialect.segment.sqlserver.hint.OptionHintSegment;
 import org.apache.shardingsphere.test.it.sql.parser.internal.asserts.SQLCaseAssertContext;
 import org.apache.shardingsphere.test.it.sql.parser.internal.asserts.segment.SQLSegmentAssert;
 import org.apache.shardingsphere.test.it.sql.parser.internal.asserts.segment.limit.LimitClauseAssert;
@@ -34,6 +35,8 @@
 
 import java.util.Optional;
 
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -46,7 +49,7 @@ public final class UpdateStatementAssert {
 
     /**
      * Assert update statement is correct with expected parser result.
-     * 
+     *
      * @param assertContext assert context
      * @param actual actual update statement
      * @param expected expected parser result
@@ -57,6 +60,7 @@ public static void assertIs(final SQLCaseAssertContext assertContext, final Upda
         assertWhereClause(assertContext, actual, expected);
         assertOrderByClause(assertContext, actual, expected);
         assertLimitClause(assertContext, actual, expected);
+        assertOptionHint(assertContext, actual, expected);
     }
 
     private static void assertTable(final SQLCaseAssertContext assertContext, final UpdateStatement actual, final UpdateStatementTestCase expected) {
@@ -100,4 +104,15 @@ private static void assertLimitClause(final SQLCaseAssertContext assertContext,
             SQLSegmentAssert.assertIs(assertContext, limitSegment.get(), expected.getLimitClause());
         }
     }
+
+    private static void assertOptionHint(final SQLCaseAssertContext assertContext, final UpdateStatement actual, final UpdateStatementTestCase expected) {
+        Optional<OptionHintSegment> optionHintSegment = UpdateStatementHandler.getOptionHintSegment(actual);
+        if (null == expected.getOptionHint()) {
+            assertFalse(optionHintSegment.isPresent(), assertContext.getText("Actual option hint segment should not exist."));
+        } else {
+            assertTrue(optionHintSegment.isPresent(), assertContext.getText("Actual option hint segment should exist."));
+            assertThat(assertContext.getText("Option hint text assertion error: "), optionHintSegment.get().getText(), is(expected.getOptionHint().getText()));
+            SQLSegmentAssert.assertIs(assertContext, optionHintSegment.get(), expected.getOptionHint());
+        }
+    }
 }
diff --git a/...e/test/it/sql/parser/internal/cases/parser/jaxb/segment/impl/hint/ExpectedOptionHint.java b/...e/test/it/sql/parser/internal/cases/parser/jaxb/segment/impl/hint/ExpectedOptionHint.java
@@ -0,0 +1,35 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.shardingsphere.test.it.sql.parser.internal.cases.parser.jaxb.segment.impl.hint;
+
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import org.apache.shardingsphere.test.it.sql.parser.internal.cases.parser.jaxb.segment.AbstractExpectedSQLSegment;
+
+import javax.xml.bind.annotation.XmlAttribute;
+
+/**
+ * Expected option hint.
+ **/
+@RequiredArgsConstructor
+@Getter
+public class ExpectedOptionHint extends AbstractExpectedSQLSegment {
+
+    @XmlAttribute
+    private String text;
+}
diff --git a/.../test/it/sql/parser/internal/cases/parser/jaxb/statement/dml/UpdateStatementTestCase.java b/.../test/it/sql/parser/internal/cases/parser/jaxb/statement/dml/UpdateStatementTestCase.java
@@ -19,6 +19,7 @@
 
 import lombok.Getter;
 import lombok.Setter;
+import org.apache.shardingsphere.test.it.sql.parser.internal.cases.parser.jaxb.segment.impl.hint.ExpectedOptionHint;
 import org.apache.shardingsphere.test.it.sql.parser.internal.cases.parser.jaxb.segment.impl.limit.ExpectedLimitClause;
 import org.apache.shardingsphere.test.it.sql.parser.internal.cases.parser.jaxb.segment.impl.orderby.ExpectedOrderByClause;
 import org.apache.shardingsphere.test.it.sql.parser.internal.cases.parser.jaxb.segment.impl.set.ExpectedSetClause;
@@ -49,4 +50,7 @@ public final class UpdateStatementTestCase extends SQLParserTestCase {
 
     @XmlElement(name = "limit")
     private ExpectedLimitClause limitClause;
+
+    @XmlElement(name = "option-hint")
+    private ExpectedOptionHint optionHint;
 }
-Original file line number
+Diff line change
@@ Expand Up / @@ -19,6 +19,7 @@ grammar SQLServerStatement; @@
     import Comments, TCLStatement, StoreProcedure, DALStatement;
     execute
         : (select
         | insert
@@ Expand Down @@