From 46433fbc4a166d8458564d2aba9f6aa09b8d9859 Mon Sep 17 00:00:00 2001
From: yydeng626 <yydeng626@gmail.com>
Date: Sat, 17 Feb 2024 15:03:35 +0800
Subject: [PATCH] Support parsing SQL Server SELECT d.name, sql #29181

---
 .../antlr4/imports/sqlserver/DMLStatement.g4  |  8 ++-
 .../statement/SQLServerStatementVisitor.java  |  3 +
 .../dml/SQLServerSelectStatement.java         | 12 ++++
 .../src/main/resources/case/dml/insert.xml    | 59 +++++++++++++++++++
 .../resources/sql/supported/dml/insert.xml    |  1 +
 5 files changed, 81 insertions(+), 2 deletions(-)

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
index 4df2d9a0a84c5..d93a3299d43b7 100644
--- a/parser/sql/dialect/sqlserver/src/main/antlr4/imports/sqlserver/DMLStatement.g4
+++ b/parser/sql/dialect/sqlserver/src/main/antlr4/imports/sqlserver/DMLStatement.g4
@@ -40,7 +40,7 @@ insertExecClause
     ;
 
 withTableHint
-    : WITH LP_ (tableHintLimited+) RP_
+    : WITH? LP_ (tableHintLimited+) RP_
     ;
 
 exec
@@ -98,7 +98,7 @@ aggregationClause
     ;
 
 selectClause
-    : selectWithClause? SELECT duplicateSpecification? projections intoClause? fromClause? whereClause? groupByClause? havingClause? orderByClause? forClause?
+    : selectWithClause? SELECT duplicateSpecification? projections intoClause? (fromClause withTempTable? withTableHint?)? whereClause? groupByClause? havingClause? orderByClause? forClause?
     ;
 
 duplicateSpecification
@@ -178,6 +178,10 @@ subquery
     : LP_ aggregationClause RP_
     ;
 
+withTempTable
+    : WITH LP_ (columnName dataType) (COMMA_ columnName dataType)* RP_
+    ;
+
 withClause
     : WITH cteClauseSet
     ;
diff --git a/parser/sql/dialect/sqlserver/src/main/java/org/apache/shardingsphere/sql/parser/sqlserver/visitor/statement/SQLServerStatementVisitor.java b/parser/sql/dialect/sqlserver/src/main/java/org/apache/shardingsphere/sql/parser/sqlserver/visitor/statement/SQLServerStatementVisitor.java
index 0aebdcc885852..d86e3676d0812 100644
--- a/parser/sql/dialect/sqlserver/src/main/java/org/apache/shardingsphere/sql/parser/sqlserver/visitor/statement/SQLServerStatementVisitor.java
+++ b/parser/sql/dialect/sqlserver/src/main/java/org/apache/shardingsphere/sql/parser/sqlserver/visitor/statement/SQLServerStatementVisitor.java
@@ -938,6 +938,9 @@ public ASTNode visitSelectClause(final SelectClauseContext ctx) {
             TableSegment tableSource = (TableSegment) visit(ctx.fromClause().tableReferences());
             result.setFrom(tableSource);
         }
+        if (null != ctx.withTableHint()) {
+            result.setWithTableHintSegment((WithTableHintSegment) visit(ctx.withTableHint()));
+        }
         if (null != ctx.whereClause()) {
             result.setWhere((WhereSegment) visit(ctx.whereClause()));
         }
diff --git a/parser/sql/statement/src/main/java/org/apache/shardingsphere/sql/parser/sql/dialect/statement/sqlserver/dml/SQLServerSelectStatement.java b/parser/sql/statement/src/main/java/org/apache/shardingsphere/sql/parser/sql/dialect/statement/sqlserver/dml/SQLServerSelectStatement.java
index 2800c7980f29a..5e7f16e441db7 100644
--- a/parser/sql/statement/src/main/java/org/apache/shardingsphere/sql/parser/sql/dialect/statement/sqlserver/dml/SQLServerSelectStatement.java
+++ b/parser/sql/statement/src/main/java/org/apache/shardingsphere/sql/parser/sql/dialect/statement/sqlserver/dml/SQLServerSelectStatement.java
@@ -22,6 +22,7 @@
 import org.apache.shardingsphere.sql.parser.sql.common.segment.generic.WithSegment;
 import org.apache.shardingsphere.sql.parser.sql.common.segment.generic.table.TableSegment;
 import org.apache.shardingsphere.sql.parser.sql.common.statement.dml.SelectStatement;
+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;
@@ -38,6 +39,8 @@ public final class SQLServerSelectStatement extends SelectStatement implements S
     
     private TableSegment intoSegment;
     
+    private WithTableHintSegment withTableHintSegment;
+    
     /**
      * Get order by segment.
      *
@@ -64,4 +67,13 @@ public Optional<WithSegment> getWithSegment() {
     public Optional<TableSegment> getIntoSegment() {
         return Optional.ofNullable(intoSegment);
     }
+    
+    /**
+     * Get with table hint segment.
+     *
+     * @return with table hint segment.
+     */
+    public Optional<WithTableHintSegment> getWithTableHintSegment() {
+        return Optional.ofNullable(withTableHintSegment);
+    }
 }
diff --git a/test/it/parser/src/main/resources/case/dml/insert.xml b/test/it/parser/src/main/resources/case/dml/insert.xml
index d12092c498e36..918aa58877b38 100644
--- a/test/it/parser/src/main/resources/case/dml/insert.xml
+++ b/test/it/parser/src/main/resources/case/dml/insert.xml
@@ -3711,4 +3711,63 @@
             </from>
         </select>
     </insert>
+    
+    <insert sql-case-id="insert_with_select_with_table_hint_no_with_keyword">
+        <table name="#tmpdbs" start-index="12" stop-index="18" />
+        <columns start-index="20" stop-index="47">
+            <column name="dbid" start-index="21" stop-index="26" start-delimiter="[" end-delimiter="]" />
+            <column name="dbname" start-index="29" stop-index="36" start-delimiter="[" end-delimiter="]" />
+            <column name="isdone" start-index="39" stop-index="46" start-delimiter="[" end-delimiter="]" />
+        </columns>
+        <select>
+            <projections start-index="56" stop-index="75">
+                <column-projection name="database_id" start-index="56" stop-index="66" />
+                <column-projection name="name" start-index="69" stop-index="72" />
+                <expression-projection text="0" start-index="75" stop-index="75">
+                    <expr>
+                        <literal-expression value="0" start-index="75" stop-index="75" />
+                    </expr>
+                </expression-projection>
+            </projections>
+            <from>
+                <simple-table name="databases" start-index="82" stop-index="101">
+                    <owner name="sys" start-index="89" stop-index="91">
+                        <owner name="master" start-index="82" stop-index="87" />
+                    </owner>
+                </simple-table>
+            </from>
+            <table-hints start-index="103" stop-index="110">
+                <table-hint value="NOLOCK" start-index="104" stop-index="109" />
+            </table-hints>
+            <where start-index="112" stop-index="147">
+                <expr>
+                    <binary-operation-expression start-index="118" stop-index="147">
+                        <left>
+                            <binary-operation-expression start-index="118" stop-index="133">
+                                <left>
+                                    <column name="is_read_only" start-index="118" stop-index="129" />
+                                </left>
+                                <right>
+                                    <literal-expression value="0" start-index="133" stop-index="133" />
+                                </right>
+                                <operator>=</operator>
+                            </binary-operation-expression>
+                        </left>
+                        <right>
+                            <binary-operation-expression start-index="139" stop-index="147">
+                                <left>
+                                    <column name="state" start-index="139" stop-index="143" />
+                                </left>
+                                <right>
+                                    <literal-expression value="0" start-index="147" stop-index="147" />
+                                </right>
+                                <operator>=</operator>
+                            </binary-operation-expression>
+                        </right>
+                        <operator>AND</operator>
+                    </binary-operation-expression>
+                </expr>
+            </where>
+        </select>
+    </insert>
 </sql-parser-test-cases>
diff --git a/test/it/parser/src/main/resources/sql/supported/dml/insert.xml b/test/it/parser/src/main/resources/sql/supported/dml/insert.xml
index d230ffc6f6b9c..1812613d93608 100644
--- a/test/it/parser/src/main/resources/sql/supported/dml/insert.xml
+++ b/test/it/parser/src/main/resources/sql/supported/dml/insert.xml
@@ -141,4 +141,5 @@
     <sql-case id="insert_with_currency_value_2" value="INSERT INTO [HR].[Employees]([SSN], [FirstName], [LastName], [Salary]) VALUES ('990-00-6818', N'Kim', N'Abercrombie', $55415)" db-types="SQLServer"/>
     <sql-case id="insert_with_cross_database_values" value="INSERT [SourceDatabase].[dbo].[SourceTable] VALUES (1, N'Bob'),(2, N'Susan')" db-types="SQLServer"/>
     <sql-case id="insert_with_cross_database_select" value="INSERT [DestinationDatabase].[dbo].[DestTable_InMem] SELECT * FROM [SourceDatabase].[dbo].[SourceTable]" db-types="SQLServer"/>
+    <sql-case id="insert_with_select_with_table_hint_no_with_keyword" value="INSERT INTO #tmpdbs ([dbid], [dbname], [isdone]) SELECT database_id, name, 0 FROM master.sys.databases (NOLOCK) WHERE is_read_only = 0 AND state = 0" db-types="SQLServer"/>
 </sql-cases>