For update prototype 2

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"
@@ -448,10 +452,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,52 @@ 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;
+}
+{
+    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>
+            {
+                waitSeconds = Long.parseLong(token.image);
+            }
+        |
+            <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/exec/ExecutionServiceImpl.java

@@ -17,6 +17,7 @@
 
 package org.apache.ignite.internal.processors.query.calcite.exec;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
@@ -29,11 +30,13 @@
 import org.apache.calcite.plan.Contexts;
 import org.apache.calcite.plan.RelOptUtil;
 import org.apache.calcite.rel.core.TableModify;
+import org.apache.calcite.schema.SchemaPlus;
 import org.apache.calcite.sql.SqlInsert;
 import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.tools.FrameworkConfig;
 import org.apache.ignite.IgniteCheckedException;
 import org.apache.ignite.IgniteException;
+import org.apache.ignite.cache.CacheEntry;
 import org.apache.ignite.cache.query.FieldsQueryCursor;
 import org.apache.ignite.cache.query.QueryCancelledException;
 import org.apache.ignite.calcite.CalciteQueryEngineConfiguration;
@@ -45,16 +48,21 @@
 import org.apache.ignite.internal.managers.eventstorage.DiscoveryEventListener;
 import org.apache.ignite.internal.managers.eventstorage.GridEventStorageManager;
 import org.apache.ignite.internal.processors.affinity.AffinityTopologyVersion;
+import org.apache.ignite.internal.processors.cache.CacheEntryImplEx;
 import org.apache.ignite.internal.processors.cache.CacheObjectUtils;
 import org.apache.ignite.internal.processors.cache.CacheObjectValueContext;
+import org.apache.ignite.internal.processors.cache.GridCacheContext;
 import org.apache.ignite.internal.processors.cache.GridCachePartitionExchangeManager;
+import org.apache.ignite.internal.processors.cache.IgniteInternalCache;
 import org.apache.ignite.internal.processors.cache.QueryCursorImpl;
 import org.apache.ignite.internal.processors.cache.distributed.near.GridNearTxLocal;
 import org.apache.ignite.internal.processors.cache.query.CacheQueryType;
 import org.apache.ignite.internal.processors.cache.query.GridCacheQueryType;
 import org.apache.ignite.internal.processors.cache.query.IgniteQueryErrorCode;
+import org.apache.ignite.internal.processors.cache.version.GridCacheVersion;
 import org.apache.ignite.internal.processors.failure.FailureProcessor;
 import org.apache.ignite.internal.processors.performancestatistics.PerformanceStatisticsProcessor;
+import org.apache.ignite.internal.processors.query.GridQueryFieldMetadata;
 import org.apache.ignite.internal.processors.query.IgniteSQLException;
 import org.apache.ignite.internal.processors.query.QueryProperties;
 import org.apache.ignite.internal.processors.query.calcite.CalciteQueryProcessor;
@@ -98,19 +106,22 @@
 import org.apache.ignite.internal.processors.query.calcite.prepare.PrepareServiceImpl;
 import org.apache.ignite.internal.processors.query.calcite.prepare.QueryPlan;
 import org.apache.ignite.internal.processors.query.calcite.prepare.QueryPlanCache;
+import org.apache.ignite.internal.processors.query.calcite.prepare.SelectForUpdatePlan;
 import org.apache.ignite.internal.processors.query.calcite.prepare.ddl.CreateTableCommand;
 import org.apache.ignite.internal.processors.query.calcite.rel.IgniteIndexBound;
 import org.apache.ignite.internal.processors.query.calcite.rel.IgniteIndexCount;
 import org.apache.ignite.internal.processors.query.calcite.rel.IgniteIndexScan;
 import org.apache.ignite.internal.processors.query.calcite.rel.IgniteRel;
 import org.apache.ignite.internal.processors.query.calcite.rel.IgniteTableModify;
 import org.apache.ignite.internal.processors.query.calcite.rel.IgniteTableScan;
+import org.apache.ignite.internal.processors.query.calcite.schema.CacheTableDescriptor;
 import org.apache.ignite.internal.processors.query.calcite.schema.IgniteTable;
 import org.apache.ignite.internal.processors.query.calcite.schema.SchemaHolder;
 import org.apache.ignite.internal.processors.query.calcite.type.IgniteTypeFactory;
 import org.apache.ignite.internal.processors.query.calcite.util.AbstractService;
 import org.apache.ignite.internal.processors.query.calcite.util.Commons;
 import org.apache.ignite.internal.processors.query.calcite.util.ConvertingClosableIterator;
+import org.apache.ignite.internal.processors.query.calcite.util.IgniteResource;
 import org.apache.ignite.internal.processors.query.calcite.util.ListFieldsQueryCursor;
 import org.apache.ignite.internal.processors.query.running.HeavyQueriesTracker;
 import org.apache.ignite.internal.processors.security.SecurityUtils;
@@ -541,6 +552,9 @@ private FragmentPlan prepareFragment(BaseQueryContext ctx, String jsonFragment)
             case DDL:
                 return executeDdl(qry, (DdlPlan)plan);
 
+            case FOR_UPDATE:
+                return executeForUpdate(qry, (SelectForUpdatePlan)plan);
+
             default:
                 throw new AssertionError("Unexpected plan type: " + plan);
         }
@@ -581,6 +595,161 @@ private FieldsQueryCursor<List<?>> executeDdl(RootQuery<Row> qry, DdlPlan plan)
         }
     }
 
+    /**
+     * Executes a {@code SELECT ... FOR UPDATE} plan.
+     *
+     * <ol>
+     *   <li>Validates that the current transaction is PESSIMISTIC.</li>
+     *   <li>Runs the inner SELECT (which has {@code _KEY} appended) and materialises all rows.</li>
+     *   <li>Extracts keys from the last column of each row.</li>
+     *   <li>Calls {@code cache.getEntries(keys)} to obtain current entry versions.</li>
+     *   <li>Creates a savepoint, acquires pessimistic locks via {@code lockTxEntries()},
+     *       and releases the savepoint on success (or rolls back on failure).</li>
+     *   <li>Returns a cursor with only the user-visible columns (the appended _KEY is stripped).</li>
+     * </ol>
+     */
+    private FieldsQueryCursor<List<?>> executeForUpdate(RootQuery<Row> qry, SelectForUpdatePlan plan) {
+        GridNearTxLocal userTx = Commons.queryTransaction(qry.context(), ctx.cache().context());
+
+        if (userTx == null || !userTx.pessimistic())
+            throw new IgniteSQLException(
+                IgniteResource.INSTANCE.selectForUpdateRequiresPessimisticTx().str(),
+                IgniteQueryErrorCode.UNSUPPORTED_OPERATION);
+
+        // Run the inner SELECT (with _KEY appended) and collect all rows.
+        ListFieldsQueryCursor<?> innerCursor = mapAndExecutePlan(qry, plan.innerPlan());
+        List<List<?>> rows = innerCursor.getAll();
+
+        int userColCount = plan.userColumnCount();
+
+        if (rows.isEmpty()) {
+            // Nothing to lock – return an empty cursor with user-only field metadata.
+            QueryCursorImpl<List<?>> resCur = new QueryCursorImpl<>(Collections.emptyList(), null, false);
+
+            IgniteTypeFactory typeFactory = qry.context().typeFactory();
+            List<GridQueryFieldMetadata> meta = plan.innerPlan().fieldsMetadata().queryFieldsMetadata(typeFactory);
+
+            resCur.fieldsMeta(meta.subList(0, userColCount));
+
+            return resCur;
+        }
+
+        List<CacheEntry<Object, Object>> entries = new ArrayList<>();
+
+        for (List<?> row : rows) {
+            GridCacheVersion ver = (GridCacheVersion) row.get(row.size() - 1);
+            Object val = row.get(row.size() - 2);
+            Object key = row.get(row.size() - 3);
+
+            entries.add(new CacheEntryImplEx<>(key, val, ver));
+        }
+
+        // Resolve the target cache via the schema holder.
+        SchemaPlus schemaPlus = schemaHolder.schema(plan.schemaName());
+
+        if (schemaPlus == null)
+            throw new IgniteSQLException("Schema not found: " + plan.schemaName(),
+                IgniteQueryErrorCode.SCHEMA_NOT_FOUND);
+
+        IgniteTable igniteTable = (IgniteTable)schemaPlus.getTable(plan.tableName());
+
+        if (igniteTable == null)
+            throw new IgniteSQLException("Table not found: " + plan.tableName(),
+                IgniteQueryErrorCode.TABLE_NOT_FOUND);
+
+        GridCacheContext<Object, Object> cctx =
+            (GridCacheContext<Object, Object>)((CacheTableDescriptor)igniteTable.descriptor()).cacheContext();
+
+        IgniteInternalCache<Object, Object> cache = cctx.cache().keepBinary();
+
+        try {
+            // Compute lock wait timeout.
+            // waitSeconds: null = use tx remaining time (0), 0 = NOWAIT (-1), positive = ms.
+            Long waitSeconds = plan.waitSeconds();
+            long waitMs;
+
+            if (waitSeconds == null)
+                waitMs = 0L;
+            else if (waitSeconds == 0L)
+                waitMs = -1L;
+            else
+                waitMs = waitSeconds * 1000L;
+
+            // lockTxEntries() requires the transaction to be bound to the current thread
+            // (it checks cctx.tm().threadLocalTx()). Resume it here and suspend afterwards,
+            // following the same pattern as ModifyNode.invokeInsideTransaction().
+            userTx.resume();
+
+            try {
+                // Create a savepoint so that a failed lock attempt can be rolled back without aborting the whole tx.
+                String spName = "_for_update_" + UUID.randomUUID();
+
+                userTx.savepoint(spName, false);
+
+                boolean locked = false;
+
+                try {
+                    locked = cache.lockTxEntries(entries, waitMs);
+                }
+                catch (IgniteCheckedException e) {
+                    try {
+                        userTx.rollbackToSavepoint(spName);
+                    }
+                    catch (Exception rollbackEx) {
+                        e.addSuppressed(rollbackEx);
+                    }
+
+                    throw new IgniteSQLException("Failed to acquire locks for SELECT FOR UPDATE",
+                        IgniteQueryErrorCode.CONCURRENT_UPDATE, e);
+                }
+
+                if (!locked) {
+                    try {
+                        userTx.rollbackToSavepoint(spName);
+                    }
+                    catch (IgniteCheckedException rollbackEx) {
+                        throw new IgniteSQLException("Failed to rollback savepoint after lock failure",
+                            IgniteQueryErrorCode.UNKNOWN, rollbackEx);
+                    }
+
+                    throw new IgniteSQLException(
+                        IgniteResource.INSTANCE.selectForUpdateLockFailed().str(),
+                        IgniteQueryErrorCode.CONCURRENT_UPDATE);
+                }
+
+                try {
+                    userTx.releaseSavepoint(spName);
+                }
+                catch (IgniteCheckedException e) {
+                    throw new IgniteSQLException("Failed to release savepoint after successful lock",
+                        IgniteQueryErrorCode.UNKNOWN, e);
+                }
+            }
+            finally {
+                userTx.suspend();
+            }
+        }
+        catch (IgniteCheckedException e) {
+            throw new IgniteSQLException("Failed to get cache entries for SELECT FOR UPDATE",
+                IgniteQueryErrorCode.UNKNOWN, e);
+        }
+
+        // Return user rows with the appended _KEY column stripped.
+        List<List<?>> userRows = new ArrayList<>(rows.size());
+
+        for (List<?> row : rows)
+            userRows.add(row.subList(0, userColCount));
+
+        QueryCursorImpl<List<?>> resCur = new QueryCursorImpl<>(userRows, null, false);
+
+        IgniteTypeFactory typeFactory = qry.context().typeFactory();
+        List<GridQueryFieldMetadata> meta = plan.innerPlan().fieldsMetadata().queryFieldsMetadata(typeFactory);
+
+        resCur.fieldsMeta(meta.subList(0, userColCount));
+
+        return resCur;
+    }
+
     /**
      * Checks that query is initiated by UDF.
      *

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

@@ -365,7 +365,8 @@ private static boolean isAsCall(SqlNode node) {
         }
 
         CalciteCatalogReader catalogReader = this.catalogReader.withSchemaPath(schemaPath);
-        SqlValidator validator = new IgniteSqlValidator(operatorTbl, catalogReader, typeFactory, validatorCfg, ctx.parameters());
+        IgniteSqlValidator validator = new IgniteSqlValidator(operatorTbl, catalogReader, typeFactory, validatorCfg, ctx.parameters());
+        validator.allowTechnicalColumns(ctx.allowTechnicalColumns());
         SqlToRelConverter sqlToRelConverter = sqlToRelConverter(validator, catalogReader, sqlToRelConverterCfg);
         RelRoot root = sqlToRelConverter.convertQuery(sqlNode, true, false);
         root = root.withRel(sqlToRelConverter.decorrelate(sqlNode, root.rel));
@@ -425,8 +426,13 @@ public String dump() {
 
     /** */
     private SqlValidator validator() {
-        if (validator == null)
-            validator = new IgniteSqlValidator(operatorTbl, catalogReader, typeFactory, validatorCfg, ctx.parameters());
+        if (validator == null) {
+            IgniteSqlValidator v = new IgniteSqlValidator(operatorTbl, catalogReader, typeFactory, validatorCfg, ctx.parameters());
+
+            v.allowTechnicalColumns(ctx.allowTechnicalColumns());
+
+            validator = v;
+        }
 
         return validator;
     }