[feature][dingo-calcite,dingo-exec] Support select... for update

dingo-calcite/src/main/codegen/templates/Parser.jj

@@ -1434,7 +1434,7 @@ SqlNode SqlSelect() :
             new SqlNodeList(selectList, Span.of(selectList).pos()),
             fromClause, where, groupBy, having, windowDecls, null, null, null,
             new SqlNodeList(hints, getPos()),
-            exportOptions, flashbackQuery, flashBackStr, flashBackTsoStr);
+            exportOptions, flashbackQuery, flashBackStr, flashBackTsoStr, forUpdate);
     }
 }
 

dingo-calcite/src/main/java/io/dingodb/calcite/DingoParser.java

@@ -51,12 +51,14 @@
 import io.dingodb.calcite.grammar.dql.SqlBackUpTimePoint;
 import io.dingodb.calcite.grammar.dql.SqlBackUpTsoPoint;
 import io.dingodb.calcite.grammar.dql.SqlNextAutoIncrement;
+import io.dingodb.calcite.grammar.dql.SqlSelect;
 import io.dingodb.calcite.grammar.dql.SqlShow;
 import io.dingodb.calcite.grammar.dql.SqlStartGc;
 import io.dingodb.calcite.meta.DingoRelMetadataProvider;
 import io.dingodb.calcite.program.DecorrelateProgram;
 import io.dingodb.calcite.rel.DingoCost;
 import io.dingodb.calcite.rel.LogicalExportData;
+import io.dingodb.calcite.rel.LogicalForUpdate;
 import io.dingodb.calcite.rel.logical.LogicalDingoRoot;
 import io.dingodb.calcite.rule.DingoRules;
 import io.dingodb.calcite.runtime.DingoResource;
@@ -82,6 +84,7 @@
 import org.apache.calcite.plan.ConventionTraitDef;
 import org.apache.calcite.plan.RelOptCluster;
 import org.apache.calcite.plan.RelOptRule;
+import org.apache.calcite.plan.RelOptTable;
 import org.apache.calcite.plan.RelTraitSet;
 import org.apache.calcite.plan.ViewExpanders;
 import org.apache.calcite.plan.volcano.AbstractConverter;
@@ -266,7 +269,7 @@ public RelRoot convert(@NonNull SqlNode sqlNode, boolean needsValidation) {
             );
 
             if (needExport(sqlNode)) {
-                io.dingodb.calcite.grammar.dql.SqlSelect sqlSelect = (io.dingodb.calcite.grammar.dql.SqlSelect) sqlNode;
+                SqlSelect sqlSelect = (SqlSelect) sqlNode;
                 validatorExportParam(sqlSelect.getExportOptions());
                 relNode = new LogicalExportData(
                     cluster,
@@ -283,22 +286,32 @@ public RelRoot convert(@NonNull SqlNode sqlNode, boolean needsValidation) {
                     context.getTimeZone()
                 );
             }
+            if (forUpdate(sqlNode)) {
+                SqlSelect sqlSelect = (SqlSelect) sqlNode;
+                relNode = new LogicalForUpdate(
+                    cluster,
+                    planner.emptyTraitSet(),
+                    relRoot.rel,
+                    Objects.requireNonNull(
+                        sqlValidator.getNamespace(Objects.requireNonNull(sqlSelect.getFrom()))).getTable().unwrap(RelOptTable.class)
+                );
+            }
         }
         // Insert a `DingoRoot` to collect the results.
         return relRoot.withRel(new LogicalDingoRoot(cluster, planner.emptyTraitSet(), relNode, selection));
     }
 
     public static boolean needExport(@NonNull SqlNode sqlNode) {
-        if (sqlNode instanceof io.dingodb.calcite.grammar.dql.SqlSelect) {
-            io.dingodb.calcite.grammar.dql.SqlSelect sqlSelect = (io.dingodb.calcite.grammar.dql.SqlSelect) sqlNode;
+        if (sqlNode instanceof SqlSelect) {
+            SqlSelect sqlSelect = (SqlSelect) sqlNode;
             return sqlSelect.isExport();
         }
         return false;
     }
 
     public static boolean flashBackQuery(@NonNull SqlNode sqlNode) {
-        if (sqlNode instanceof io.dingodb.calcite.grammar.dql.SqlSelect) {
-            io.dingodb.calcite.grammar.dql.SqlSelect sqlSelect = (io.dingodb.calcite.grammar.dql.SqlSelect) sqlNode;
+        if (sqlNode instanceof SqlSelect) {
+            SqlSelect sqlSelect = (SqlSelect) sqlNode;
             if (sqlSelect.getFrom() instanceof FlashBackSqlIdentifier) {
                 return true;
             } else {
@@ -308,6 +321,14 @@ public static boolean flashBackQuery(@NonNull SqlNode sqlNode) {
         return false;
     }
 
+    public static boolean forUpdate(@NonNull SqlNode sqlNode) {
+        if (sqlNode instanceof SqlSelect) {
+            SqlSelect sqlSelect = (SqlSelect) sqlNode;
+            return sqlSelect.isForUpdate();
+        }
+        return false;
+    }
+
     /**
      * Optimize a {@link RelNode} tree.
      *

dingo-calcite/src/main/java/io/dingodb/calcite/grammar/dql/SqlSelect.java

@@ -39,6 +39,8 @@ public class SqlSelect extends org.apache.calcite.sql.SqlSelect {
 
     long flashBackTso;
 
+    boolean forUpdate;
+
     public SqlSelect(SqlParserPos pos,
                      @Nullable SqlNodeList keywordList,
                      SqlNodeList selectList,
@@ -69,14 +71,16 @@ public SqlSelect(SqlParserPos pos,
                      ExportOptions exportOptions,
                      boolean flashbackQuery,
                      String flashBackStr,
-                     String flashBackTsoStr) {
+                     String flashBackTsoStr,
+                     boolean forUpdate) {
         super(pos, keywordList, selectList, from, where, groupBy, having, windowDecls, orderBy, offset, fetch, hints);
         this.exportOptions =  exportOptions;
         this.flashBackQuery = flashbackQuery;
         this.flashBackTimeStr = flashBackStr;
         if (flashBackTsoStr != null) {
             this.flashBackTso = Long.parseLong(flashBackTsoStr);
         }
+        this.forUpdate = forUpdate;
     }
 
     public SqlSelect(SqlParserPos pos,

dingo-calcite/src/main/java/io/dingodb/calcite/rel/DingoForUpdate.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2021 DataCanvas
+ *
+ * Licensed 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 io.dingodb.calcite.rel;
+
+import io.dingodb.calcite.visitor.DingoRelVisitor;
+import lombok.Getter;
+import org.apache.calcite.plan.RelOptCluster;
+import org.apache.calcite.plan.RelOptTable;
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.SingleRel;
+import org.checkerframework.checker.nullness.qual.NonNull;
+
+import java.util.List;
+
+public class DingoForUpdate extends SingleRel implements DingoRel {
+
+    @Getter
+    private RelOptTable table;
+
+    /**
+     * Creates a <code>SingleRel</code>.
+     *
+     * @param cluster Cluster this relational expression belongs to
+     * @param traits
+     * @param input   Input relational expression
+     */
+    public DingoForUpdate(RelOptCluster cluster, RelTraitSet traits, RelNode input, RelOptTable table) {
+        super(cluster, traits, input);
+        this.table = table;
+    }
+
+    @Override
+    public <T> T accept(@NonNull DingoRelVisitor<T> visitor) {
+        return visitor.visit(this);
+    }
+
+    @Override
+    public RelNode copy(RelTraitSet traitSet, List<RelNode> inputs) {
+        return new DingoForUpdate(getCluster(), traitSet, sole(inputs), table);
+    }
+}

dingo-calcite/src/main/java/io/dingodb/calcite/rel/LogicalForUpdate.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2021 DataCanvas
+ *
+ * Licensed 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 io.dingodb.calcite.rel;
+
+import io.dingodb.calcite.DingoRelOptTable;
+import lombok.Getter;
+import org.apache.calcite.plan.RelOptCluster;
+import org.apache.calcite.plan.RelOptTable;
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.SingleRel;
+
+import java.util.List;
+
+public class LogicalForUpdate extends SingleRel {
+
+    @Getter
+    private final RelOptTable table;
+
+    /**
+     * Creates a <code>SingleRel</code>.
+     *
+     * @param cluster Cluster this relational expression belongs to
+     * @param traits
+     * @param input   Input relational expression
+     */
+    public LogicalForUpdate(RelOptCluster cluster, RelTraitSet traits, RelNode input, RelOptTable table) {
+        super(cluster, traits, input);
+        this.table = table;
+    }
+
+    @Override
+    public RelNode copy(RelTraitSet traitSet, List<RelNode> inputs) {
+        return new LogicalForUpdate(getCluster(), traitSet, sole(inputs), table);
+    }
+}

dingo-calcite/src/main/java/io/dingodb/calcite/rule/DingoForUpdateRule.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2021 DataCanvas
+ *
+ * Licensed 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 io.dingodb.calcite.rule;
+
+import io.dingodb.calcite.rel.DingoForUpdate;
+import io.dingodb.calcite.rel.LogicalForUpdate;
+import io.dingodb.calcite.traits.DingoConvention;
+import io.dingodb.calcite.traits.DingoRelStreaming;
+import org.apache.calcite.plan.Convention;
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.convert.ConverterRule;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+public class DingoForUpdateRule extends ConverterRule {
+
+    public static final Config DEFAULT = Config.INSTANCE
+        .withConversion(
+            LogicalForUpdate.class,
+            Convention.NONE,
+            DingoConvention.INSTANCE,
+            "DingoForUpdateRule"
+        )
+        .withRuleFactory(DingoForUpdateRule::new);
+
+    protected DingoForUpdateRule(Config config) {
+        super(config);
+    }
+
+    @Override
+    public @Nullable RelNode convert(RelNode rel) {
+        LogicalForUpdate logicalForUpdate = (LogicalForUpdate) rel;
+        RelTraitSet traits = logicalForUpdate.getTraitSet()
+            .replace(DingoConvention.INSTANCE)
+            .replace(DingoRelStreaming.ROOT);
+        return new DingoForUpdate(
+            logicalForUpdate.getCluster(),
+            traits,
+            convert(logicalForUpdate.getInput(), traits),
+            logicalForUpdate.getTable());
+    }
+}

dingo-calcite/src/main/java/io/dingodb/calcite/rule/DingoRules.java

@@ -110,6 +110,9 @@ public final class DingoRules {
     public static final DingoExportDataRule EXPORT_DATA_RULE =
         DingoExportDataRule.DEFAULT.toRule(DingoExportDataRule.class);
 
+    public static final DingoForUpdateRule FOR_UPDATE_RULE =
+        DingoForUpdateRule.DEFAULT.toRule(DingoForUpdateRule.class);
+
     public static final DingoWithoutPriModifyRule WITHOUT_PRI_DELETE_RULE =
         DingoWithoutPriModifyRule.Config.DELETE.toRule();
 
@@ -223,6 +226,7 @@ public final class DingoRules {
         FILTER_SUB_QUERY_TO_CORRELATE,
         JOIN_SUB_QUERY_TO_CORRELATE,
         EXPORT_DATA_RULE,
+        FOR_UPDATE_RULE,
         WITHOUT_PRI_DELETE_RULE,
         WITHOUT_PRI_UPDATE_RULE,
         DINGO_VECTOR_PROJECT_RULE,

dingo-calcite/src/main/java/io/dingodb/calcite/visitor/DingoExplainVisitor.java

@@ -26,6 +26,7 @@
 import io.dingodb.calcite.rel.DingoDocument;
 import io.dingodb.calcite.rel.DingoExportData;
 import io.dingodb.calcite.rel.DingoFilter;
+import io.dingodb.calcite.rel.DingoForUpdate;
 import io.dingodb.calcite.rel.DingoFunctionScan;
 import io.dingodb.calcite.rel.DingoGetByIndex;
 import io.dingodb.calcite.rel.DingoGetByIndexMerge;
@@ -394,6 +395,16 @@ public Explain visit(@NonNull DingoExportData rel) {
         return explain1;
     }
 
+    @Override
+    public Explain visit(@NonNull DingoForUpdate rel) {
+        String filter = "";
+        String tableNames = "";
+        return new Explain(
+            "dingoForUpdate", rel.getRowCount(), "root",
+            tableNames, filter
+        );
+    }
+
     @Override
     public Explain visit(@NonNull IndexFullScan rel) {
         StringBuilder info = new StringBuilder();