*: support select for update of tables (pingcap#1333)

Lvnszn · Sep 13, 2021 · 530bf89 · 530bf89
1 parent d551970
commit 530bf89
Show file tree

Hide file tree

Showing 4 changed files with 6,448 additions and 6,307 deletions.
diff --git a/ast/dml.go b/ast/dml.go
@@ -615,6 +615,7 @@ const (
 type SelectLockInfo struct {
 	LockType SelectLockType
 	WaitSec  uint64
+	Tables   []*TableName
 }
 
 // String implements fmt.Stringer.
@@ -1319,11 +1320,48 @@ func (n *SelectStmt) Restore(ctx *format.RestoreCtx) error {
 		ctx.WritePlain(" ")
 		switch n.LockInfo.LockType {
 		case SelectLockNone:
+		case SelectLockForUpdateNoWait:
+			ctx.WriteKeyWord("for update")
+			if len(n.LockInfo.Tables) != 0 {
+				ctx.WriteKeyWord(" OF ")
+				restoreTables(ctx, n.LockInfo.Tables)
+			}
+			ctx.WriteKeyWord(" nowait")
 		case SelectLockForUpdateWaitN:
-			ctx.WriteKeyWord(n.LockInfo.LockType.String())
+			ctx.WriteKeyWord("for update")
+			if len(n.LockInfo.Tables) != 0 {
+				ctx.WriteKeyWord(" OF ")
+				restoreTables(ctx, n.LockInfo.Tables)
+			}
+			ctx.WriteKeyWord(" wait")
 			ctx.WritePlainf(" %d", n.LockInfo.WaitSec)
+		case SelectLockForShareNoWait:
+			ctx.WriteKeyWord("for share")
+			if len(n.LockInfo.Tables) != 0 {
+				ctx.WriteKeyWord(" OF ")
+				restoreTables(ctx, n.LockInfo.Tables)
+			}
+			ctx.WriteKeyWord(" nowait")
+		case SelectLockForUpdateSkipLocked:
+			ctx.WriteKeyWord("for update")
+			if len(n.LockInfo.Tables) != 0 {
+				ctx.WriteKeyWord(" OF ")
+				restoreTables(ctx, n.LockInfo.Tables)
+			}
+			ctx.WriteKeyWord(" skip locked")
+		case SelectLockForShareSkipLocked:
+			ctx.WriteKeyWord("for share")
+			if len(n.LockInfo.Tables) != 0 {
+				ctx.WriteKeyWord(" OF ")
+				restoreTables(ctx, n.LockInfo.Tables)
+			}
+			ctx.WriteKeyWord(" skip locked")
 		default:
 			ctx.WriteKeyWord(n.LockInfo.LockType.String())
+			if len(n.LockInfo.Tables) != 0 {
+				ctx.WriteKeyWord(" OF ")
+				restoreTables(ctx, n.LockInfo.Tables)
+			}
 		}
 	}
 
@@ -1336,6 +1374,18 @@ func (n *SelectStmt) Restore(ctx *format.RestoreCtx) error {
 	return nil
 }
 
+func restoreTables(ctx *format.RestoreCtx, ts []*TableName) error {
+	for i, v := range ts {
+		if err := v.Restore(ctx); err != nil {
+			return errors.Annotate(err, "An error occurred while restore SelectStmt.LockInfo")
+		}
+		if i != len(ts)-1 {
+			ctx.WritePlain(", ")
+		}
+	}
+	return nil
+}
+
 // Accept implements Node Accept interface.
 func (n *SelectStmt) Accept(v Visitor) (Node, bool) {
 	newNode, skipChildren := v.Enter(n)
@@ -1437,6 +1487,16 @@ func (n *SelectStmt) Accept(v Visitor) (Node, bool) {
 		n.Limit = node.(*Limit)
 	}
 
+	if n.LockInfo != nil {
+		for i, t := range n.LockInfo.Tables {
+			node, ok := t.Accept(v)
+			if !ok {
+				return n, false
+			}
+			n.LockInfo.Tables[i] = node.(*TableName)
+		}
+	}
+
 	return v.Leave(n)
 }