planner: support table aliases in FOR UPDATE OF Clause (#65532)

cryo-zd · web-flow · commit bab4993b6fd6 · 2026-03-18T22:25:14.000Z
close #63035
diff --git a/pkg/planner/core/logical_plan_builder.go b/pkg/planner/core/logical_plan_builder.go
@@ -4090,12 +4090,21 @@ func (b *PlanBuilder) buildSelect(ctx context.Context, sel *ast.SelectStmt) (p b
 				// Besides, it will only lock the metioned in `of` part.
 				b.ctx.GetSessionVars().StmtCtx.LockTableIDs[tNameW.TableInfo.ID] = struct{}{}
 			}
-			dbName := getLowerDB(tName.Schema, b.ctx.GetSessionVars())
+			// Use the already-resolved DBInfo to derive the privilege-check DB name.
+			// For OF-alias targets, tName.Schema is empty; falling back to currentDB via getLowerDB would
+			// authorize against the wrong database when the aliased table lives in a different schema.
+			var dbName string
+			if tNameW.DBInfo != nil {
+				dbName = tNameW.DBInfo.Name.L
+			} else {
+				dbName = getLowerDB(tName.Schema, b.ctx.GetSessionVars())
+			}
+
 			var authErr error
 			if user := b.ctx.GetSessionVars().User; user != nil {
-				authErr = plannererrors.ErrTableaccessDenied.GenWithStackByArgs("SELECT with locking clause", user.AuthUsername, user.AuthHostname, tNameW.Name.L)
+				authErr = plannererrors.ErrTableaccessDenied.GenWithStackByArgs("SELECT with locking clause", user.AuthUsername, user.AuthHostname, tNameW.TableInfo.Name.L)
 			}
-			b.visitInfo = appendVisitInfo(b.visitInfo, mysql.DeletePriv|mysql.UpdatePriv|mysql.LockTablesPriv, dbName, tNameW.Name.L, "", authErr)
+			b.visitInfo = appendVisitInfo(b.visitInfo, mysql.DeletePriv|mysql.UpdatePriv|mysql.LockTablesPriv, dbName, tNameW.TableInfo.Name.L, "", authErr)
 		}
 		p, err = b.buildSelectLock(p, l)
 		if err != nil {
diff --git a/pkg/planner/core/preprocess.go b/pkg/planner/core/preprocess.go
@@ -132,6 +132,7 @@ func Preprocess(ctx context.Context, sctx sessionctx.Context, node *resolve.Node
 		sctx:               sctx,
 		tableAliasInJoin:   make([]map[string]any, 0),
 		preprocessWith:     &preprocessWith{cteCanUsed: make([]string, 0), cteBeforeOffset: make([]int, 0)},
+		lockSelectCtxStack: make([]lockSelectCtx, 0),
 		staleReadProcessor: staleread.NewStaleReadProcessor(ctx, sctx),
 		varsMutable:        make(map[string]struct{}),
 		varsReadonly:       make(map[string]struct{}),
@@ -238,6 +239,11 @@ type preprocessor struct {
 	tableAliasInJoin []map[string]any
 	preprocessWith   *preprocessWith
 
+	// lockSelectCtxStack tracks lock-clause resolution state for each SELECT. Each level keeps both
+	// the LockInfo TableName nodes that should skip normal preprocess resolution and the FROM-clause
+	// table references collected during traversal for later lock-target binding.
+	lockSelectCtxStack []lockSelectCtx
+
 	staleReadProcessor staleread.Processor
 
 	varsMutable  map[string]struct{}
@@ -262,6 +268,9 @@ func (p *preprocessor) Enter(in ast.Node) (out ast.Node, skipChildren bool) {
 			p.preprocessWith.cteStack = append(p.preprocessWith.cteStack, node.With.CTEs)
 		}
 		p.checkSelectNoopFuncs(node)
+		// SelectStmt.Accept visits FROM before LockInfo, so one per-SELECT context can collect FROM
+		// table refs during traversal and later bind LockInfo targets without re-walking the FROM tree.
+		p.pushLockSelectCtx(node)
 	case *ast.SetOprStmt:
 		if node.With != nil {
 			p.preprocessWith.cteStack = append(p.preprocessWith.cteStack, node.With.CTEs)
@@ -655,7 +664,19 @@ func (p *preprocessor) Leave(in ast.Node) (out ast.Node, ok bool) {
 			p.err = plannererrors.ErrUnknownExplainFormat.GenWithStackByArgs(x.Format)
 		}
 	case *ast.TableName:
+		// Skip resolving TableName nodes that come from locking clauses; they are bound later against FROM aliases.
+		if lockCtx := p.getLockSelectCtxStackTop(); lockCtx != nil {
+			if _, ok := lockCtx.lockClauseTables[x]; ok {
+				break
+			}
+		}
 		p.handleTableName(x)
+	case *ast.TableSource:
+		if lockCtx := p.getLockSelectCtxStackTop(); lockCtx != nil {
+			if v, ok := x.Source.(*ast.TableName); ok {
+				lockCtx.collect(v, x.AsName)
+			}
+		}
 	case *ast.Join:
 		if len(p.tableAliasInJoin) > 0 {
 			p.tableAliasInJoin = p.tableAliasInJoin[:len(p.tableAliasInJoin)-1]
@@ -716,6 +737,10 @@ func (p *preprocessor) Leave(in ast.Node) (out ast.Node, ok bool) {
 		if x.With != nil {
 			p.preprocessWith.cteStack = p.preprocessWith.cteStack[0 : len(p.preprocessWith.cteStack)-1]
 		}
+		if lockCtx := p.getLockSelectCtxStackTop(); lockCtx != nil {
+			p.checkLockClauseTables(x, lockCtx)
+			p.popLockSelectCtx()
+		}
 	case *ast.SetOprStmt:
 		if x.With != nil {
 			p.preprocessWith.cteStack = p.preprocessWith.cteStack[0 : len(p.preprocessWith.cteStack)-1]
@@ -1800,6 +1825,174 @@ func (p *preprocessor) handleTableName(tn *ast.TableName) {
 	})
 }
 
+// lockRef records the FROM table and its alias (if any) for lock-clause name resolution.
+type lockRef struct {
+	table *ast.TableName
+	alias ast.CIStr
+}
+
+type lockSelectCtx struct {
+	// lockClauseTables are the TableName nodes that syntactically belong to this SELECT's locking
+	// clause. They skip handleTableName because they are resolved later against the FROM clause.
+	lockClauseTables map[*ast.TableName]struct{}
+	aliasMap         map[string]lockRef
+	qualifiedMap     map[string]lockRef
+	// orderedRefs preserves the left-to-right FROM order for unqualified OF fallback.
+	// For example, `FROM db1.t, db2.t` records `db1.t` before `db2.t`, so fallback resolution for `FOR UPDATE OF t` will try `db1.t` first.
+	// When both aliased and unaliased entries share the same base name, checkLockClauseTables first
+	// searches orderedRefs for an exact unaliased match, and only then falls back to aliased entries
+	// for backward compatibility.
+	orderedRefs []lockRef
+}
+
+func newLockSelectCtx(lockTables []*ast.TableName) lockSelectCtx {
+	lockClauseTables := make(map[*ast.TableName]struct{}, len(lockTables))
+	for _, tn := range lockTables {
+		lockClauseTables[tn] = struct{}{}
+	}
+	return lockSelectCtx{
+		lockClauseTables: lockClauseTables,
+		aliasMap:         make(map[string]lockRef),
+		qualifiedMap:     make(map[string]lockRef),
+		orderedRefs:      make([]lockRef, 0),
+	}
+}
+
+func (c *lockSelectCtx) collect(tableName *ast.TableName, asName ast.CIStr) {
+	ref := lockRef{
+		table: tableName,
+		alias: asName,
+	}
+	if asName.L != "" {
+		c.aliasMap[asName.L] = ref
+	}
+	if tableName.Schema.L != "" {
+		qualifiedKey := tableName.Schema.L + "." + tableName.Name.L
+		// Prefer an exact unaliased `schema.table` entry. If only aliased references exist in FROM,
+		// keep one as the backward-compatibility fallback for `OF schema.table`.
+		if asName.L == "" || c.qualifiedMap[qualifiedKey].table == nil {
+			c.qualifiedMap[qualifiedKey] = ref
+		}
+		if asName.L != "" {
+			// Keep backward compatibility for `OF schema.table` while also supporting `OF schema.alias`.
+			c.qualifiedMap[tableName.Schema.L+"."+asName.L] = ref
+		}
+	}
+	c.orderedRefs = append(c.orderedRefs, ref)
+}
+
+func (p *preprocessor) pushLockSelectCtx(sel *ast.SelectStmt) {
+	var lockTables []*ast.TableName
+	if sel.LockInfo != nil {
+		lockTables = sel.LockInfo.Tables
+	}
+	p.lockSelectCtxStack = append(p.lockSelectCtxStack, newLockSelectCtx(lockTables))
+}
+
+func (p *preprocessor) getLockSelectCtxStackTop() *lockSelectCtx {
+	if len(p.lockSelectCtxStack) == 0 {
+		return nil
+	}
+	return &p.lockSelectCtxStack[len(p.lockSelectCtxStack)-1]
+}
+
+func (p *preprocessor) popLockSelectCtx() {
+	if len(p.lockSelectCtxStack) == 0 {
+		return
+	}
+	p.lockSelectCtxStack = p.lockSelectCtxStack[:len(p.lockSelectCtxStack)-1]
+}
+
+// checkLockClauseTables validates/attaches tables referenced by SELECT ... FOR UPDATE/LOCK IN SHARE MODE.
+// It binds lock targets to the tables in FROM and keeps backward compatibility by accepting
+// base-table names even when aliases exist, with a warning to guide users to explicit aliases.
+func (p *preprocessor) checkLockClauseTables(stmt *ast.SelectStmt, lockCtx *lockSelectCtx) {
+	if stmt.LockInfo == nil || stmt.LockInfo.LockType == ast.SelectLockNone || len(stmt.LockInfo.Tables) == 0 {
+		return
+	}
+	if lockCtx == nil {
+		emptyCtx := newLockSelectCtx(nil)
+		lockCtx = &emptyCtx
+	}
+	warnedCompatRefs := make(map[string]struct{})
+
+	for _, ref := range stmt.LockInfo.Tables {
+		name := ref.Name.L
+		var matched lockRef
+		var matchedByAlias bool
+
+		if ref.Schema.L != "" {
+			name = ref.Schema.L + "." + ref.Name.L
+			matched = lockCtx.qualifiedMap[name]
+			matchedByAlias = matched.alias.L != "" && ref.Name.L == matched.alias.L
+			if matchedByAlias {
+				ref.IsAlias = true
+			}
+		} else if t, ok := lockCtx.aliasMap[name]; ok {
+			ref.IsAlias = true
+			matched = t
+			matchedByAlias = true
+		}
+
+		// If still not found and the lock clause is unqualified, first try an exact unaliased match.
+		// For example, `SELECT * FROM db.t AS a, db.t FOR UPDATE OF t` should resolve `t` to the
+		// unaliased `db.t`, not to `a`. Only if no unaliased entry exists do we fall back to aliased
+		// entries for backward compatibility. `SELECT * FROM db1.t, db2.t FOR UPDATE OF t` still
+		// resolves `t` to `db1.t` because lockCtx.orderedRefs preserves the left-to-right FROM order.
+		if matched.table == nil && ref.Schema.L == "" {
+			for _, tblRef := range lockCtx.orderedRefs {
+				if tblRef.alias.L == "" && tblRef.table.Name.L == name {
+					matched = tblRef
+					break
+				}
+			}
+			if matched.table == nil {
+				for _, tblRef := range lockCtx.orderedRefs {
+					if tblRef.alias.L != "" && tblRef.table.Name.L == name {
+						matched = tblRef
+						break
+					}
+				}
+			}
+		}
+
+		if matched.table == nil {
+			if ref.Schema.L != "" {
+				name = ref.Schema.O + "." + ref.Name.O
+			} else {
+				name = ref.Name.O
+			}
+			p.err = plannererrors.ErrUnknownTable.GenWithStackByArgs(name, "locking clause")
+			break
+		}
+
+		// Keep backward compatibility: if an aliased table is referenced by its base name
+		// in OF, accept it and emit a warning to guide users to the alias form.
+		if !matchedByAlias && matched.alias.L != "" {
+			warnKey := ref.Schema.L + "." + ref.Name.L + "->" + matched.alias.L
+			if _, seen := warnedCompatRefs[warnKey]; !seen {
+				p.sctx.GetSessionVars().StmtCtx.AppendWarning(errors.NewNoStackErrorf(
+					"FOR UPDATE OF references the base table name while the table is aliased. Use the alias '%s' in OF to make the lock target explicit.",
+					matched.alias.O,
+				))
+				warnedCompatRefs[warnKey] = struct{}{}
+			}
+		}
+
+		tNameW := p.resolveCtx.GetTableName(matched.table)
+		if tNameW == nil {
+			// CTE has no *model.HintedTable, we need to skip it.
+			continue
+		}
+
+		p.resolveCtx.AddTableName(&resolve.TableNameW{
+			TableName: ref,
+			DBInfo:    tNameW.DBInfo,
+			TableInfo: tNameW.TableInfo,
+		})
+	}
+}
+
 func (p *preprocessor) checkNotInRepair(tn *ast.TableName) {
 	tableInfo, dbInfo := domainutil.RepairInfo.GetRepairedTableInfoByTableName(tn.Schema.L, tn.Name.L)
 	if dbInfo == nil {
diff --git a/tests/integrationtest/r/planner/core/issuetest/planner_issue.result b/tests/integrationtest/r/planner/core/issuetest/planner_issue.result
@@ -1482,3 +1482,68 @@ select * from t1 natural join t2 order by t1.a;
 a
 1
 drop table t1, t2;
+set @old_db := database();
+drop database if exists lock_db1;
+drop database if exists lock_db2;
+create database lock_db1;
+create database lock_db2;
+use lock_db1;
+create table t(id int primary key, v int);
+insert into t values (1, 100);
+select ue1_0.id from t ue1_0 where ue1_0.id=1 for update of t;
+id
+1
+Level	Code	Message
+Warning	1105	FOR UPDATE OF references the base table name while the table is aliased. Use the alias 'ue1_0' in OF to make the lock target explicit.
+select ue1_0.id from t ue1_0 where ue1_0.id=1 for update of ue1_0;
+id
+1
+select * from t ue1_0 for update of t;
+id	v
+1	100
+Level	Code	Message
+Warning	1105	FOR UPDATE OF references the base table name while the table is aliased. Use the alias 'ue1_0' in OF to make the lock target explicit.
+select * from t ue1_0 for update of t, t, ue1_0;
+id	v
+1	100
+Level	Code	Message
+Warning	1105	FOR UPDATE OF references the base table name while the table is aliased. Use the alias 'ue1_0' in OF to make the lock target explicit.
+select * from t ue1_0 for update of lock_db1.t;
+id	v
+1	100
+Level	Code	Message
+Warning	1105	FOR UPDATE OF references the base table name while the table is aliased. Use the alias 'ue1_0' in OF to make the lock target explicit.
+select * from t ue1_0 for update of lock_db1.ue1_0;
+id	v
+1	100
+select * from t ue1_0, t where ue1_0.id = t.id for update of t;
+id	v	id	v
+1	100	1	100
+select * from t, t ue1_0 where ue1_0.id = t.id for update of lock_db1.t;
+id	v	id	v
+1	100	1	100
+select * from t ue1_0 for update of non_exist_table;
+Error 1109 (42S02): Unknown table 'non_exist_table' in locking clause
+use lock_db2;
+create table t(id int primary key, v int);
+insert into t values (1, 200);
+select * from lock_db1.t ue1_0 for update of t;
+id	v
+1	100
+Level	Code	Message
+Warning	1105	FOR UPDATE OF references the base table name while the table is aliased. Use the alias 'ue1_0' in OF to make the lock target explicit.
+select * from lock_db1.t ue1_0, lock_db2.t ue1_1 for update of t;
+id	v	id	v
+1	100	1	200
+Level	Code	Message
+Warning	1105	FOR UPDATE OF references the base table name while the table is aliased. Use the alias 'ue1_0' in OF to make the lock target explicit.
+with cte_t as (select * from t) select * from cte_t for update of cte_t;
+id	v
+1	200
+set @old_db = ifnull(@old_db, 'test');
+set @stmt = concat('use ', @old_db);
+prepare _ldbstmt from @stmt;
+execute _ldbstmt;
+drop database if exists lock_db1;
+drop database if exists lock_db2;
+deallocate prepare _ldbstmt;
diff --git a/tests/integrationtest/t/planner/core/issuetest/planner_issue.test b/tests/integrationtest/t/planner/core/issuetest/planner_issue.test
@@ -914,3 +914,42 @@ insert into t1 values (1);
 insert into t2 values (1);
 select * from t1 natural join t2 order by t1.a;
 drop table t1, t2;
+
+# Issue63035
+set @old_db := database();
+drop database if exists lock_db1;
+drop database if exists lock_db2;
+create database lock_db1;
+create database lock_db2;
+
+use lock_db1;
+create table t(id int primary key, v int);
+insert into t values (1, 100);
+--enable_warnings
+select ue1_0.id from t ue1_0 where ue1_0.id=1 for update of t;
+select ue1_0.id from t ue1_0 where ue1_0.id=1 for update of ue1_0;
+select * from t ue1_0 for update of t;
+select * from t ue1_0 for update of t, t, ue1_0;
+select * from t ue1_0 for update of lock_db1.t;
+select * from t ue1_0 for update of lock_db1.ue1_0;
+select * from t ue1_0, t where ue1_0.id = t.id for update of t;
+select * from t, t ue1_0 where ue1_0.id = t.id for update of lock_db1.t;
+
+--error 1109
+select * from t ue1_0 for update of non_exist_table;
+
+use lock_db2;
+create table t(id int primary key, v int);
+insert into t values (1, 200);
+select * from lock_db1.t ue1_0 for update of t;
+select * from lock_db1.t ue1_0, lock_db2.t ue1_1 for update of t;
+with cte_t as (select * from t) select * from cte_t for update of cte_t;
+--disable_warnings
+
+set @old_db = ifnull(@old_db, 'test');
+set @stmt = concat('use ', @old_db);
+prepare _ldbstmt from @stmt;
+execute _ldbstmt;
+drop database if exists lock_db1;
+drop database if exists lock_db2;
+deallocate prepare _ldbstmt;
