Disallow Insert with Duplicate key update and Replace Into queries on foreign key column, set locks on fk queries

go/vt/sqlparser/ast.go

@@ -63,6 +63,7 @@ type (
 		GetOrderBy() OrderBy
 		GetLimit() *Limit
 		SetLimit(*Limit)
+		GetLock() Lock
 		SetLock(lock Lock)
 		SetInto(into *SelectInto)
 		SetWith(with *With)

go/vt/sqlparser/ast_funcs.go

@@ -1050,6 +1050,11 @@ func (node *Select) GetLimit() *Limit {
 	return node.Limit
 }
 
+// GetLock returns the lock clause
+func (node *Select) GetLock() Lock {
+	return node.Lock
+}
+
 // SetLock sets the lock clause
 func (node *Select) SetLock(lock Lock) {
 	node.Lock = lock
@@ -1196,6 +1201,11 @@ func (node *Union) GetColumns() SelectExprs {
 	return node.Left.GetColumns()
 }
 
+// GetLock returns the lock clause
+func (node *Union) GetLock() Lock {
+	return node.Lock
+}
+
 // SetLock sets the lock clause
 func (node *Union) SetLock(lock Lock) {
 	node.Lock = lock

go/vt/sqlparser/comments.go

@@ -273,6 +273,13 @@ func (c *ParsedComments) Length() int {
 	return len(c.comments)
 }
 
+func (c *ParsedComments) GetComments() Comments {
+	if c != nil {
+		return c.comments
+	}
+	return nil
+}
+
 func (c *ParsedComments) Prepend(comment string) Comments {
 	if c == nil {
 		return Comments{comment}

go/vt/vtgate/planbuilder/operator_transformers.go

@@ -439,6 +439,12 @@ func transformRoutePlan(ctx *plancontext.PlanningContext, op *operators.Route) (
 
 	switch stmt := stmt.(type) {
 	case sqlparser.SelectStatement:
+		if op.Lock != sqlparser.NoLock {
+			stmt.SetLock(op.Lock)
+		}
+		if op.Comments != nil {
+			stmt.SetComments(op.Comments.GetComments())
+		}
 		return buildRouteLogicalPlan(ctx, op, stmt)
 	case *sqlparser.Update:
 		return buildUpdateLogicalPlan(ctx, op, dmlOp, stmt)

go/vt/vtgate/planbuilder/operators/ast2op.go

@@ -280,12 +280,13 @@ func addColumnEquality(ctx *plancontext.PlanningContext, expr sqlparser.Expr) {
 // createSelectionOp creates the selection operator to select the parent columns for the foreign key constraints.
 // The Select statement looks something like this - `SELECT <parent_columns_in_fk for all the foreign key constraints> FROM <parent_table> WHERE <where_clause_of_update>`
 // TODO (@Harshit, @GuptaManan100): Compress the columns in the SELECT statement, if there are multiple foreign key constraints using the same columns.
-func createSelectionOp(ctx *plancontext.PlanningContext, selectExprs []sqlparser.SelectExpr, tableExprs sqlparser.TableExprs, where *sqlparser.Where, limit *sqlparser.Limit) (ops.Operator, error) {
+func createSelectionOp(ctx *plancontext.PlanningContext, selectExprs []sqlparser.SelectExpr, tableExprs sqlparser.TableExprs, where *sqlparser.Where, limit *sqlparser.Limit, lock sqlparser.Lock) (ops.Operator, error) {
 	selectionStmt := &sqlparser.Select{
 		SelectExprs: selectExprs,
 		From:        tableExprs,
 		Where:       where,
 		Limit:       limit,
+		Lock:        lock,
 	}
 	// There are no foreign keys to check for a select query, so we can pass anything for verifyAllFKs and fkToIgnore.
 	return createOpFromStmt(ctx, selectionStmt, false /* verifyAllFKs */, "" /* fkToIgnore */)

go/vt/vtgate/planbuilder/operators/delete.go

@@ -198,7 +198,7 @@ func createFkCascadeOpForDelete(ctx *plancontext.PlanningContext, parentOp ops.O
 		}
 		fkChildren = append(fkChildren, fkChild)
 	}
-	selectionOp, err := createSelectionOp(ctx, selectExprs, delStmt.TableExprs, delStmt.Where, nil)
+	selectionOp, err := createSelectionOp(ctx, selectExprs, delStmt.TableExprs, delStmt.Where, nil, sqlparser.ForUpdateLock)
 	if err != nil {
 		return nil, err
 	}

go/vt/vtgate/planbuilder/operators/operator.go

@@ -89,6 +89,21 @@ func PlanQuery(ctx *plancontext.PlanningContext, stmt sqlparser.Statement) (ops.
 		return nil, ctx.SemTable.NotSingleRouteErr
 	}
 
+	// set lock and comments on the route to be set on the sql query on conversion.
+	_ = rewrite.Visit(op, func(op ops.Operator) error {
+		route, ok := op.(*Route)
+		if !ok {
+			return nil
+		}
+		if stmtWithComments, ok := stmt.(sqlparser.Commented); ok {
+			route.Comments = stmtWithComments.GetParsedComments()
+		}
+		if stmtWithLock, ok := stmt.(sqlparser.SelectStatement); ok {
+			route.Lock = stmtWithLock.GetLock()
+		}
+		return nil
+	})
+
 	return op, err
 }
 

go/vt/vtgate/planbuilder/operators/route.go

@@ -42,6 +42,9 @@ type (
 
 		Ordering []RouteOrdering
 
+		Comments *sqlparser.ParsedComments
+		Lock     sqlparser.Lock
+
 		ResultColumns int
 	}
 

go/vt/vtgate/planbuilder/operators/update.go

@@ -306,7 +306,7 @@ func createFKCascadeOp(ctx *plancontext.PlanningContext, parentOp ops.Operator,
 		fkChildren = append(fkChildren, fkChild)
 	}
 
-	selectionOp, err := createSelectionOp(ctx, selectExprs, updStmt.TableExprs, updStmt.Where, nil)
+	selectionOp, err := createSelectionOp(ctx, selectExprs, updStmt.TableExprs, updStmt.Where, nil, sqlparser.ForUpdateLock)
 	if err != nil {
 		return nil, err
 	}
@@ -551,7 +551,8 @@ func createFkVerifyOpForParentFKForUpdate(ctx *plancontext.PlanningContext, updS
 				sqlparser.NewJoinCondition(joinCond, nil)),
 		},
 		sqlparser.NewWhere(sqlparser.WhereClause, whereCond),
-		sqlparser.NewLimitWithoutOffset(1))
+		sqlparser.NewLimitWithoutOffset(1),
+		sqlparser.ShareModeLock)
 }
 
 // Each child foreign key constraint is verified by a join query of the form:
@@ -604,5 +605,6 @@ func createFkVerifyOpForChildFKForUpdate(ctx *plancontext.PlanningContext, updSt
 				sqlparser.NewJoinCondition(joinCond, nil)),
 		},
 		sqlparser.NewWhere(sqlparser.WhereClause, whereCond),
-		sqlparser.NewLimitWithoutOffset(1))
+		sqlparser.NewLimitWithoutOffset(1),
+		sqlparser.ShareModeLock)
 }

go/vt/vtgate/planbuilder/select.go

@@ -234,12 +234,6 @@ func newBuildSelectPlan(
 
 	optimizePlan(plan)
 
-	if sel, isSel := selStmt.(*sqlparser.Select); isSel {
-		if err = setMiscFunc(plan, sel); err != nil {
-			return nil, nil, err
-		}
-	}
-
 	if err = plan.Wireup(ctx); err != nil {
 		return nil, nil, err
 	}
@@ -396,38 +390,6 @@ func shouldRetryAfterPredicateRewriting(plan logicalPlan) bool {
 		len(sysTableTableSchema) == 0
 }
 
-func setMiscFunc(in logicalPlan, sel *sqlparser.Select) error {
-	_, err := visit(in, func(plan logicalPlan) (bool, logicalPlan, error) {
-		switch node := plan.(type) {
-		case *route:
-			err := copyCommentsAndLocks(node.Select, sel, node.eroute.Opcode)
-			if err != nil {
-				return false, nil, err
-			}
-			return true, node, nil
-		}
-		return true, plan, nil
-	})
-
-	if err != nil {
-		return err
-	}
-	return nil
-}
-
-func copyCommentsAndLocks(statement sqlparser.SelectStatement, sel *sqlparser.Select, opcode engine.Opcode) error {
-	query := sqlparser.GetFirstSelect(statement)
-	query.Comments = sel.Comments
-	query.Lock = sel.Lock
-	if sel.Into != nil {
-		if opcode != engine.Unsharded {
-			return vterrors.VT12001("INTO on sharded keyspace")
-		}
-		query.Into = sel.Into
-	}
-	return nil
-}
-
 func handleDualSelects(sel *sqlparser.Select, vschema plancontext.VSchema) (engine.Primitive, error) {
 	if !isOnlyDual(sel) {
 		return nil, nil

go/vt/vtgate/planbuilder/testdata/foreignkey_cases.json

@@ -677,7 +677,7 @@
                   "Sharded": false
                 },
                 "FieldQuery": "select col2 from u_tbl2 where 1 != 1",
-                "Query": "select col2 from u_tbl2 where (col2) in ::fkc_vals",
+                "Query": "select col2 from u_tbl2 where (col2) in ::fkc_vals for update",
                 "Table": "u_tbl2"
               },
               {
@@ -727,7 +727,7 @@
                   "Sharded": false
                 },
                 "FieldQuery": "select col9 from u_tbl9 where 1 != 1",
-                "Query": "select col9 from u_tbl9 where (col9) in ::fkc_vals2 and (col9) not in ('foo')",
+                "Query": "select col9 from u_tbl9 where (col9) in ::fkc_vals2 and (col9) not in ('foo') for update",
                 "Table": "u_tbl9"
               },
               {
@@ -893,7 +893,7 @@
                   "Sharded": false
                 },
                 "FieldQuery": "select col2 from u_tbl2 where 1 != 1",
-                "Query": "select col2 from u_tbl2 where (col2) in ::fkc_vals",
+                "Query": "select col2 from u_tbl2 where (col2) in ::fkc_vals for update",
                 "Table": "u_tbl2"
               },
               {
@@ -943,7 +943,7 @@
                   "Sharded": false
                 },
                 "FieldQuery": "select col9 from u_tbl9 where 1 != 1",
-                "Query": "select col9 from u_tbl9 where (col9) in ::fkc_vals2 and (col9) not in (2)",
+                "Query": "select col9 from u_tbl9 where (col9) in ::fkc_vals2 and (col9) not in (2) for update",
                 "Table": "u_tbl9"
               },
               {
@@ -1132,7 +1132,7 @@
                   "Sharded": false
                 },
                 "FieldQuery": "select 1 from u_tbl8 left join u_tbl9 on u_tbl9.col9 = 'foo' where 1 != 1",
-                "Query": "select 1 from u_tbl8 left join u_tbl9 on u_tbl9.col9 = 'foo' where (u_tbl8.col8) in ::fkc_vals and u_tbl9.col9 is null limit 1",
+                "Query": "select 1 from u_tbl8 left join u_tbl9 on u_tbl9.col9 = 'foo' where (u_tbl8.col8) in ::fkc_vals and u_tbl9.col9 is null limit 1 lock in share mode",
                 "Table": "u_tbl8, u_tbl9"
               },
               {
@@ -1208,7 +1208,7 @@
                   "Sharded": false
                 },
                 "FieldQuery": "select 1 from u_tbl4 left join u_tbl3 on u_tbl3.col3 = 'foo' where 1 != 1",
-                "Query": "select 1 from u_tbl4 left join u_tbl3 on u_tbl3.col3 = 'foo' where (u_tbl4.col4) in ::fkc_vals and u_tbl3.col3 is null limit 1",
+                "Query": "select 1 from u_tbl4 left join u_tbl3 on u_tbl3.col3 = 'foo' where (u_tbl4.col4) in ::fkc_vals and u_tbl3.col3 is null limit 1 lock in share mode",
                 "Table": "u_tbl3, u_tbl4"
               },
               {
@@ -1220,7 +1220,7 @@
                   "Sharded": false
                 },
                 "FieldQuery": "select 1 from u_tbl4, u_tbl9 where 1 != 1",
-                "Query": "select 1 from u_tbl4, u_tbl9 where (u_tbl4.col4) in ::fkc_vals and u_tbl4.col4 = u_tbl9.col9 limit 1",
+                "Query": "select 1 from u_tbl4, u_tbl9 where (u_tbl4.col4) in ::fkc_vals and u_tbl4.col4 = u_tbl9.col9 limit 1 lock in share mode",
                 "Table": "u_tbl4, u_tbl9"
               },
               {

go/vt/vtgate/semantics/check_invalid.go

@@ -132,6 +132,9 @@ func (a *analyzer) checkSelect(cursor *sqlparser.Cursor, node *sqlparser.Select)
 	if a.scoper.currentScope().parent != nil {
 		return &CantUseOptionHereError{Msg: errMsg}
 	}
+	if node.Into != nil {
+		return ShardedError{Inner: &UnsupportedConstruct{errString: "INTO on sharded keyspace"}}
+	}
 	return nil
 }