IGNITE-28724 Calcite engine. Support SELECT ... FOR UPDATE query

modules/calcite/src/main/codegen/config.fmpp

@@ -73,7 +73,9 @@ data: {
       "REFRESH",
       "ANALYZE",
       "MAX_CHANGED_PARTITION_ROWS_PERCENT",
-      "TOTAL"
+      "TOTAL",
+      "WAIT",
+      "NOWAIT"
     ]
 
     # List of non-reserved keywords to add;
@@ -110,6 +112,8 @@ data: {
       "ANALYZE"
       "MAX_CHANGED_PARTITION_ROWS_PERCENT"
       "TOTAL"
+      "WAIT"
+      "NOWAIT"
 
       # Keywords reserved by Calcite, but not required to be reserved in Ignite.
       "ABS"
@@ -449,10 +453,11 @@ data: {
       "SqlSavepoint()",
       "SqlRollbackToSavepoint()",
       "SqlCommitTransaction()",
-      "SqlRollbackTransaction()"
-      "SqlStatisticsAnalyze()"
-      "SqlStatisticsRefresh()"
-      "SqlStatisticsDrop()"
+      "SqlRollbackTransaction()",
+      "SqlStatisticsAnalyze()",
+      "SqlStatisticsRefresh()",
+      "SqlStatisticsDrop()",
+      "SqlSelectForUpdate()"
     ]
 
     # List of methods for parsing extensions to "CREATE [OR REPLACE]" calls.

modules/calcite/src/main/codegen/includes/parserImpls.ftl

@@ -816,3 +816,66 @@ SqlDrop SqlDropView(Span s, boolean replace) :
         return SqlDdlNodes.dropView(s.end(this), ifExists, id);
     }
 }
+
+/**
+ * Parses a query optionally followed by FOR UPDATE [OF col [, col ...]] [WAIT n | NOWAIT].
+ *
+ * When FOR UPDATE is absent the inner query node is returned unchanged, so this rule
+ * transparently handles all queries that reach the StatementParser.
+ */
+SqlNode SqlSelectForUpdate() :
+{
+    final Span s;
+    SqlNode qry;
+    SqlNodeList ofList = null;
+    List<SqlNode> ofCols = null;
+    SqlIdentifier col;
+    Long waitSeconds = null;
+    String waitValue;
+}
+{
+    qry = OrderedQueryOrExpr(ExprContext.ACCEPT_QUERY) { s = span(); }
+    [
+        LOOKAHEAD(<FOR> <UPDATE>)
+        <FOR> <UPDATE>
+        [
+            LOOKAHEAD(<OF>)
+            <OF>
+            {
+                ofCols = new ArrayList<SqlNode>();
+            }
+            col = CompoundIdentifier() { ofCols.add(col); }
+            (
+                <COMMA> col = CompoundIdentifier() { ofCols.add(col); }
+            )*
+            { ofList = new SqlNodeList(ofCols, s.pos()); }
+        ]
+        [
+            LOOKAHEAD(<WAIT>)
+            <WAIT> <UNSIGNED_INTEGER_LITERAL>
+            {
+                waitValue = token.image;
+
+                try {
+                    waitSeconds = Long.parseLong(waitValue);
+                }
+                catch (NumberFormatException ignored) {
+                    throw SqlUtil.newContextException(getPos(),
+                        IgniteResource.INSTANCE.illegalWaitTimeout(waitValue));
+                }
+
+                if (waitSeconds <= 0) {
+                    throw SqlUtil.newContextException(getPos(),
+                        IgniteResource.INSTANCE.illegalWaitTimeout(waitValue));
+                }
+            }
+        |
+            <NOWAIT>
+            {
+                waitSeconds = 0L;
+            }
+        ]
+        { return new IgniteSqlSelectForUpdate(s.end(this), qry, ofList, waitSeconds); }
+    ]
+    { return qry; }
+}

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/RootQuery.java

@@ -85,6 +85,9 @@ public class RootQuery<RowT> extends Query<RowT> implements TrackableQuery {
     /** */
     private final BaseQueryContext ctx;
 
+    /** Original query context. */
+    private final QueryContext qryCtx;
+
     /** */
     private final long plannerTimeout;
 
@@ -130,6 +133,7 @@ public RootQuery(
         );
 
         this.sql = sql;
+        this.qryCtx = qryCtx;
         this.params = params;
 
         startTs = U.currentTimeMillis();
@@ -187,6 +191,39 @@ public BaseQueryContext context() {
         return ctx;
     }
 
+    /**
+     * Creates a query for a repeated execution of the same plan.
+     *
+     * <p>The new query preserves the complete context of this query, including the user transaction,
+     * and uses only the remaining part of the original query timeout.
+     */
+    public RootQuery<RowT> retryQuery() {
+        long remainingTime = remainingTime();
+
+        if (remainingTime == 0) {
+            throw new IgniteSQLException(
+                "The query was cancelled due to timeout",
+                IgniteQueryErrorCode.QUERY_CANCELED,
+                new QueryCancelledException());
+        }
+
+        return new RootQuery<>(
+            sql,
+            ctx.schema(),
+            params,
+            QueryContext.of(cancel, qryCtx),
+            ctx.isLocal(),
+            ctx.isForcedJoinOrder(),
+            ctx.partitions(),
+            exch,
+            unregister,
+            log,
+            plannerTimeout,
+            remainingTime,
+            initiatorId
+        );
+    }
+
     /** */
     public String sql() {
         return sql;