Optimize MySQL INSERT INTO VALUES AS row_alias parse SQLStatement

parser/sql/engine/dialect/mysql/src/main/antlr4/imports/mysql/DMLStatement.g4

@@ -20,7 +20,13 @@ grammar DMLStatement;
 import BaseRule;
 
 insert
-    : INSERT insertSpecification INTO? tableName partitionNames? (insertValuesClause | setAssignmentsClause | insertSelectClause) onDuplicateKeyClause? returningClause?
+    : INSERT insertSpecification INTO? tableName partitionNames? insertBody onDuplicateKeyClause? returningClause?
+    ;
+
+insertBody
+    : insertValuesClause
+    | setAssignmentsClause valueReference?
+    | insertSelectClause valueReference?
     ;
 
 insertSpecification
@@ -47,8 +53,12 @@ insertSelectClause
     : valueReference? (LP_ fields? RP_)? (LP_ select RP_ | select)
     ;
 
+rowAlias
+    : AS alias derivedColumns?
+    ;
+
 onDuplicateKeyClause
-    : (AS identifier derivedColumns?)?  ON DUPLICATE KEY UPDATE assignment (COMMA_ assignment)*
+    : ON DUPLICATE KEY UPDATE assignment (COMMA_ assignment)*
     ;
 
 valueReference
@@ -88,7 +98,15 @@ assignment
     ;
 
 setAssignmentsClause
-    : SET assignment (COMMA_ assignment)*
+    : SET assignmentList setRowAlias?
+    ;
+
+assignmentList
+    : assignment (COMMA_ assignment)*
+    ;
+
+setRowAlias
+    : AS alias derivedColumns?
     ;
 
 assignmentValues

parser/sql/engine/dialect/mysql/src/main/java/org/apache/shardingsphere/sql/parser/engine/mysql/visitor/statement/MySQLStatementVisitor.java

@@ -73,6 +73,7 @@
 import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.HexadecimalLiteralsContext;
 import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.IdentifierContext;
 import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.IndexNameContext;
+import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.InsertBodyContext;
 import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.InsertContext;
 import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.InsertIdentifierContext;
 import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.InsertSelectClauseContext;
@@ -114,12 +115,15 @@
 import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.ReplaceContext;
 import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.ReplaceSelectClauseContext;
 import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.ReplaceValuesClauseContext;
+import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.RowAliasContext;
 import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.RowConstructorListContext;
+import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.ValueReferenceContext;
 import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.SelectContext;
 import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.SelectSpecificationContext;
 import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.SelectWithIntoContext;
 import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.SeparatorNameContext;
 import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.SetAssignmentsClauseContext;
+import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.SetRowAliasContext;
 import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.ShorthandRegularFunctionContext;
 import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.SimpleExprContext;
 import org.apache.shardingsphere.sql.parser.autogen.MySQLStatementParser.SingleTableClauseContext;
@@ -169,6 +173,8 @@
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.assignment.ColumnAssignmentSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.assignment.InsertValuesSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.assignment.SetAssignmentSegment;
+import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.assignment.RowAliasSegment;
+import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.assignment.ValueReferenceSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.column.ColumnSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.column.InsertColumnsSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.column.OnDuplicateKeyColumnsSegment;
@@ -1460,22 +1466,37 @@ public final ASTNode visitOrderByItem(final OrderByItemContext ctx) {
     @Override
     public ASTNode visitInsert(final InsertContext ctx) {
         // TODO :FIXME, since there is no segment for insertValuesClause, InsertStatement is created by sub rule.
+        InsertStatement result = (InsertStatement) visit(ctx.insertBody());
+        if (null != ctx.onDuplicateKeyClause()) {
+            result.setOnDuplicateKeyColumns((OnDuplicateKeyColumnsSegment) visit(ctx.onDuplicateKeyClause()));
+        }
+        result.setTable((SimpleTableSegment) visit(ctx.tableName()));
+        result.addParameterMarkers(getParameterMarkerSegments());
+        if (null != ctx.returningClause()) {
+            result.setReturning((ReturningSegment) visit(ctx.returningClause()));
+        }
+        return result;
+    }
+
+    @Override
+    public ASTNode visitInsertBody(final InsertBodyContext ctx) {
         InsertStatement result;
         if (null != ctx.insertValuesClause()) {
             result = (InsertStatement) visit(ctx.insertValuesClause());
         } else if (null != ctx.insertSelectClause()) {
             result = (InsertStatement) visit(ctx.insertSelectClause());
+            if (null != ctx.valueReference()) {
+                result.setValueReference((ValueReferenceSegment) visit(ctx.valueReference()));
+            }
         } else {
             result = new InsertStatement(databaseType);
             result.setSetAssignment((SetAssignmentSegment) visit(ctx.setAssignmentsClause()));
-        }
-        if (null != ctx.onDuplicateKeyClause()) {
-            result.setOnDuplicateKeyColumns((OnDuplicateKeyColumnsSegment) visit(ctx.onDuplicateKeyClause()));
-        }
-        result.setTable((SimpleTableSegment) visit(ctx.tableName()));
-        result.addParameterMarkers(getParameterMarkerSegments());
-        if (null != ctx.returningClause()) {
-            result.setReturning((ReturningSegment) visit(ctx.returningClause()));
+            if (null != ctx.valueReference()) {
+                result.setValueReference((ValueReferenceSegment) visit(ctx.valueReference()));
+            }
+            if (null == ctx.valueReference() && null != ctx.setAssignmentsClause().setRowAlias()) {
+                result.setValueReference(createValueReferenceFromSetRowAlias(ctx.setAssignmentsClause().setRowAlias()));
+            }
         }
         return result;
     }
@@ -1496,6 +1517,20 @@ public ASTNode visitInsertSelectClause(final InsertSelectClauseContext ctx) {
         return result;
     }
 
+    private ValueReferenceSegment createValueReferenceFromSetRowAlias(final SetRowAliasContext ctx) {
+        AliasSegment alias = (AliasSegment) visit(ctx.alias());
+        Collection<ColumnSegment> derivedColumns = null;
+        if (null != ctx.derivedColumns()) {
+            derivedColumns = new LinkedList<>();
+            for (AliasContext each : ctx.derivedColumns().alias()) {
+                AliasSegment aliasSegment = (AliasSegment) visit(each);
+                ColumnSegment column = new ColumnSegment(each.getStart().getStartIndex(), each.getStop().getStopIndex(), aliasSegment.getIdentifier());
+                derivedColumns.add(column);
+            }
+        }
+        return new ValueReferenceSegment(ctx.getStart().getStartIndex(), ctx.getStop().getStopIndex(), alias, derivedColumns);
+    }
+
     private SubquerySegment createInsertSelectSegment(final InsertSelectClauseContext ctx) {
         SelectStatement selectStatement = (SelectStatement) visit(ctx.select());
         selectStatement.addParameterMarkers(getParameterMarkerSegments());
@@ -1515,6 +1550,10 @@ public ASTNode visitInsertValuesClause(final InsertValuesClauseContext ctx) {
             result.setInsertColumns(new InsertColumnsSegment(ctx.start.getStartIndex() - 1, ctx.start.getStartIndex() - 1, Collections.emptyList()));
         }
         result.getValues().addAll(createInsertValuesSegments(ctx.assignmentValues()));
+        if (null != ctx.valueReference()) {
+            ValueReferenceSegment valueRef = (ValueReferenceSegment) visit(ctx.valueReference());
+            result.setValueReference(valueRef);
+        }
         return result;
     }
 
@@ -1526,6 +1565,36 @@ private Collection<InsertValuesSegment> createInsertValuesSegments(final Collect
         return result;
     }
 
+    @Override
+    public ASTNode visitRowAlias(final RowAliasContext ctx) {
+        AliasSegment alias = (AliasSegment) visit(ctx.alias());
+        Collection<ColumnSegment> derivedColumns = null;
+        if (null != ctx.derivedColumns()) {
+            derivedColumns = new LinkedList<>();
+            for (AliasContext each : ctx.derivedColumns().alias()) {
+                AliasSegment aliasSegment = (AliasSegment) visit(each);
+                ColumnSegment column = new ColumnSegment(each.getStart().getStartIndex(), each.getStop().getStopIndex(), aliasSegment.getIdentifier());
+                derivedColumns.add(column);
+            }
+        }
+        return new RowAliasSegment(ctx.getStart().getStartIndex(), ctx.getStop().getStopIndex(), alias, derivedColumns);
+    }
+
+    @Override
+    public ASTNode visitValueReference(final ValueReferenceContext ctx) {
+        AliasSegment alias = (AliasSegment) visit(ctx.alias());
+        Collection<ColumnSegment> derivedColumns = null;
+        if (null != ctx.derivedColumns()) {
+            derivedColumns = new LinkedList<>();
+            for (AliasContext each : ctx.derivedColumns().alias()) {
+                AliasSegment aliasSegment = (AliasSegment) visit(each);
+                ColumnSegment column = new ColumnSegment(each.getStart().getStartIndex(), each.getStop().getStopIndex(), aliasSegment.getIdentifier());
+                derivedColumns.add(column);
+            }
+        }
+        return new ValueReferenceSegment(ctx.getStart().getStartIndex(), ctx.getStop().getStopIndex(), alias, derivedColumns);
+    }
+
     @Override
     public ASTNode visitOnDuplicateKeyClause(final OnDuplicateKeyClauseContext ctx) {
         Collection<ColumnAssignmentSegment> columns = new LinkedList<>();
@@ -1625,7 +1694,7 @@ public ASTNode visitUpdate(final UpdateContext ctx) {
     @Override
     public ASTNode visitSetAssignmentsClause(final SetAssignmentsClauseContext ctx) {
         Collection<ColumnAssignmentSegment> assignments = new LinkedList<>();
-        for (AssignmentContext each : ctx.assignment()) {
+        for (AssignmentContext each : ctx.assignmentList().assignment()) {
             assignments.add((ColumnAssignmentSegment) visit(each));
         }
         return new SetAssignmentSegment(ctx.getStart().getStartIndex(), ctx.getStop().getStopIndex(), assignments);

parser/sql/statement/core/src/main/java/org/apache/shardingsphere/sql/parser/statement/core/segment/dml/assignment/RowAliasSegment.java

@@ -0,0 +1,52 @@
+/*
+ * 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.statement.core.segment.dml.assignment;
+
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import org.apache.shardingsphere.sql.parser.statement.core.segment.SQLSegment;
+import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.column.ColumnSegment;
+import org.apache.shardingsphere.sql.parser.statement.core.segment.generic.AliasSegment;
+
+import java.util.Collection;
+import java.util.Optional;
+
+/**
+ * Row alias segment for INSERT INTO VALUES AS row_alias syntax.
+ */
+@RequiredArgsConstructor
+@Getter
+public final class RowAliasSegment implements SQLSegment {
+
+    private final int startIndex;
+
+    private final int stopIndex;
+
+    private final AliasSegment alias;
+
+    private final Collection<ColumnSegment> derivedColumns;
+
+    /**
+     * Get derived columns.
+     *
+     * @return derived columns
+     */
+    public Optional<Collection<ColumnSegment>> getDerivedColumns() {
+        return Optional.ofNullable(derivedColumns);
+    }
+}

parser/sql/statement/core/src/main/java/org/apache/shardingsphere/sql/parser/statement/core/segment/dml/assignment/ValueReferenceSegment.java

@@ -0,0 +1,52 @@
+/*
+ * 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.statement.core.segment.dml.assignment;
+
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import org.apache.shardingsphere.sql.parser.statement.core.segment.SQLSegment;
+import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.column.ColumnSegment;
+import org.apache.shardingsphere.sql.parser.statement.core.segment.generic.AliasSegment;
+
+import java.util.Collection;
+import java.util.Optional;
+
+/**
+ * Value reference segment for INSERT INTO VALUES AS alias syntax.
+ */
+@RequiredArgsConstructor
+@Getter
+public final class ValueReferenceSegment implements SQLSegment {
+
+    private final int startIndex;
+
+    private final int stopIndex;
+
+    private final AliasSegment alias;
+
+    private final Collection<ColumnSegment> derivedColumns;
+
+    /**
+     * Get derived columns.
+     *
+     * @return derived columns
+     */
+    public Optional<Collection<ColumnSegment>> getDerivedColumns() {
+        return Optional.ofNullable(derivedColumns);
+    }
+}

parser/sql/statement/core/src/main/java/org/apache/shardingsphere/sql/parser/statement/core/statement/type/dml/InsertStatement.java

@@ -23,6 +23,7 @@
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.ReturningSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.assignment.InsertValuesSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.assignment.SetAssignmentSegment;
+import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.assignment.ValueReferenceSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.column.ColumnSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.column.InsertColumnsSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.column.OnDuplicateKeyColumnsSegment;
@@ -62,6 +63,8 @@ public final class InsertStatement extends DMLStatement {
 
     private OnDuplicateKeyColumnsSegment onDuplicateKeyColumns;
 
+    private ValueReferenceSegment valueReference;
+
     private ReturningSegment returning;
 
     private OutputSegment output;
@@ -135,6 +138,15 @@ public Optional<OnDuplicateKeyColumnsSegment> getOnDuplicateKeyColumns() {
         return Optional.ofNullable(onDuplicateKeyColumns);
     }
 
+    /**
+     * Get value reference.
+     *
+     * @return value reference
+     */
+    public Optional<ValueReferenceSegment> getValueReference() {
+        return Optional.ofNullable(valueReference);
+    }
+
     /**
      * Get set assignment.
      *