support for update in Merge statements

apache · Oct 17, 2023 · 24c440d · 24c440d
1 parent bf6f6c7
commit 24c440d
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 4 deletions.
diff --git a/...dingsphere/sqlfederation/optimizer/converter/statement/merge/MergeStatementConverter.java b/...dingsphere/sqlfederation/optimizer/converter/statement/merge/MergeStatementConverter.java
@@ -19,11 +19,22 @@
 
 import org.apache.calcite.sql.SqlMerge;
 import org.apache.calcite.sql.SqlNode;
+import org.apache.calcite.sql.SqlNodeList;
+import org.apache.calcite.sql.SqlUpdate;
 import org.apache.calcite.sql.parser.SqlParserPos;
+import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.assignment.AssignmentSegment;
+import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.column.ColumnSegment;
+import org.apache.shardingsphere.sql.parser.sql.common.segment.dml.expr.ExpressionSegment;
 import org.apache.shardingsphere.sql.parser.sql.common.statement.dml.MergeStatement;
-import org.apache.shardingsphere.sqlfederation.optimizer.converter.segment.expression.ExpressionConverter;
-import org.apache.shardingsphere.sqlfederation.optimizer.converter.segment.from.TableConverter;
-import org.apache.shardingsphere.sqlfederation.optimizer.converter.statement.SQLStatementConverter;
+import org.apache.shardingsphere.sql.parser.sql.common.statement.dml.UpdateStatement;
+import org.apache.shardingsphere.sqlfederation.compiler.converter.segment.expression.ExpressionConverter;
+import org.apache.shardingsphere.sqlfederation.compiler.converter.segment.expression.impl.ColumnConverter;
+import org.apache.shardingsphere.sqlfederation.compiler.converter.segment.from.TableConverter;
+import org.apache.shardingsphere.sqlfederation.compiler.converter.segment.where.WhereConverter;
+import org.apache.shardingsphere.sqlfederation.compiler.converter.statement.SQLStatementConverter;
+
+import java.util.List;
+import java.util.stream.Collectors;
 
 /**
  * Merge statement converter.
@@ -35,6 +46,30 @@ public SqlNode convert(final MergeStatement mergeStatement) {
         SqlNode targetTable = new TableConverter().convert(mergeStatement.getTarget()).orElseThrow(IllegalStateException::new);
         SqlNode condition = new ExpressionConverter().convert(mergeStatement.getExpression().getExpr()).get();
         SqlNode sourceTable = new TableConverter().convert(mergeStatement.getSource()).orElseThrow(IllegalStateException::new);
-        return new SqlMerge(SqlParserPos.ZERO, targetTable, condition, sourceTable, null, null, null, null);
+        SqlUpdate sqlUpdate = null;
+        if (null != mergeStatement.getUpdate()) {
+            sqlUpdate = convertUpdate(mergeStatement.getUpdate());
+        }
+        return new SqlMerge(SqlParserPos.ZERO, targetTable, condition, sourceTable, sqlUpdate, null, null, null);
+    }
+
+    private SqlUpdate convertUpdate(final UpdateStatement updateStatement) {
+        SqlNode table = new TableConverter().convert(updateStatement.getTable()).orElse(SqlNodeList.EMPTY);
+        SqlNode condition = updateStatement.getWhere().flatMap(optional -> new WhereConverter().convert(optional)).orElse(null);
+        SqlNodeList columns = new SqlNodeList(SqlParserPos.ZERO);
+        SqlNodeList expressions = new SqlNodeList(SqlParserPos.ZERO);
+        for (AssignmentSegment each : updateStatement.getAssignmentSegment().orElseThrow(IllegalStateException::new).getAssignments()) {
+            columns.addAll(convertColumn(each.getColumns()));
+            expressions.add(convertExpression(each.getValue()));
+        }
+        return new SqlUpdate(SqlParserPos.ZERO, table, columns, expressions, condition, null, null);
+    }
+
+    private List<SqlNode> convertColumn(final List<ColumnSegment> columnSegments) {
+        return columnSegments.stream().map(each -> new ColumnConverter().convert(each).orElseThrow(IllegalStateException::new)).collect(Collectors.toList());
+    }
+
+    private SqlNode convertExpression(final ExpressionSegment expressionSegment) {
+        return new ExpressionConverter().convert(expressionSegment).orElseThrow(IllegalStateException::new);
     }
 }
diff --git a/test/it/optimizer/src/test/resources/converter/merge.xml b/test/it/optimizer/src/test/resources/converter/merge.xml
@@ -19,4 +19,5 @@
 <sql-node-converter-test-cases>
     <test-cases sql-case-id="merge_into_table_using_table" expected-sql="MERGE INTO &quot;people_target&quot; USING &quot;people_source&quot; ON &quot;people_target&quot;.&quot;person_id&quot; = &quot;people_source&quot;.&quot;person_id&quot;" db-types="Oracle" sql-case-types="LITERAL" />
     <test-cases sql-case-id="merge_into_table_using_subquery_alias" expected-sql="MERGE INTO &quot;bonuses&quot; &quot;D&quot; USING (SELECT &quot;employee_id&quot;, &quot;salary&quot;, &quot;department_id&quot; FROM &quot;employees&quot; WHERE &quot;department_id&quot; = 80) &quot;S&quot; ON &quot;D&quot;.&quot;employee_id&quot; = &quot;S&quot;.&quot;employee_id&quot;" db-types="Oracle" sql-case-types="LITERAL" />
+    <test-cases sql-case-id="merge_update_table" expected-sql="MERGE INTO &quot;people_target&quot; &quot;pt&quot; USING &quot;people_source&quot; &quot;ps&quot; ON &quot;pt&quot;.&quot;person_id&quot; = &quot;ps&quot;.&quot;person_id&quot; WHEN MATCHED THEN UPDATE SET &quot;pt&quot;.&quot;first_name&quot; = &quot;ps&quot;.&quot;first_name&quot;, &quot;pt&quot;.&quot;last_name&quot; = &quot;ps&quot;.&quot;last_name&quot;, &quot;pt&quot;.&quot;title&quot; = &quot;ps&quot;.&quot;title&quot;" db-types="Oracle" />
 </sql-node-converter-test-cases>