sqlsmith: add DELETE FROM ... USING and UPDATE ... FROM support

This commit makes it so that the sqlsmith can now generate statements
of the form `DELETE FROM ... USING` and `UPDATE ... FROM`. We toss
a coin every time before deciding to add another table (meaning in 50%
cases these forms are not used, in 25% we have 1 extra table, etc). It
also adjusts the generation of the RETURNING clause for DELETE and
UPDATE to be able to pick from any of the table sources.

Release note: None

pkg/internal/sqlsmith/relational.go

@@ -875,24 +875,40 @@ func makeDelete(s *Smither) (tree.Statement, bool) {
 	return stmt, ok
 }
 
-func (s *Smither) makeDelete(refs colRefs) (*tree.Delete, *tableRef, bool) {
-	table, _, tableRef, tableRefs, ok := s.getSchemaTable()
+func (s *Smither) makeDelete(refs colRefs) (*tree.Delete, []*tableRef, bool) {
+	table, _, ref, colRefs, ok := s.getSchemaTable()
 	if !ok {
 		return nil, nil, false
 	}
+	tableRefs := []*tableRef{ref}
+
+	var hasJoinTable bool
+	var using tree.TableExprs
+	// With 50% probably add another table into the USING clause.
+	for s.coin() {
+		t, _, tableRef, cols, ok := s.getSchemaTable()
+		if !ok {
+			break
+		}
+		hasJoinTable = true
+		using = append(using, t)
+		tableRefs = append(tableRefs, tableRef)
+		colRefs = append(colRefs, cols...)
+	}
 
 	del := &tree.Delete{
 		Table:     table,
-		Where:     s.makeWhere(tableRefs, false /* hasJoinTable */),
-		OrderBy:   s.makeOrderBy(tableRefs),
+		Where:     s.makeWhere(colRefs, hasJoinTable),
+		OrderBy:   s.makeOrderBy(colRefs),
+		Using:     using,
 		Limit:     makeLimit(s),
 		Returning: &tree.NoReturningClause{},
 	}
 	if del.Limit == nil {
 		del.OrderBy = nil
 	}
 
-	return del, tableRef, true
+	return del, tableRefs, true
 }
 
 func makeDeleteReturning(s *Smither, refs colRefs, forJoin bool) (tree.TableExpr, colRefs, bool) {
@@ -919,26 +935,45 @@ func makeUpdate(s *Smither) (tree.Statement, bool) {
 	return stmt, ok
 }
 
-func (s *Smither) makeUpdate(refs colRefs) (*tree.Update, *tableRef, bool) {
-	table, _, tableRef, tableRefs, ok := s.getSchemaTable()
+func (s *Smither) makeUpdate(refs colRefs) (*tree.Update, []*tableRef, bool) {
+	table, _, ref, colRefs, ok := s.getSchemaTable()
 	if !ok {
 		return nil, nil, false
 	}
-	cols := make(map[tree.Name]*tree.ColumnTableDef)
-	for _, c := range tableRef.Columns {
-		cols[c.Name] = c
+	tableRefs := []*tableRef{ref}
+	// Each column can be set at most once. Copy colRefs to upRefs - we will
+	// remove elements from it as we use them below.
+	//
+	// Note that we need to make this copy before we append columns from other
+	// tables added into the FROM clause below.
+	upRefs := colRefs.extend()
+
+	var hasJoinTable bool
+	var from tree.TableExprs
+	// With 50% probably add another table into the FROM clause.
+	for s.coin() {
+		t, _, tableRef, cols, ok := s.getSchemaTable()
+		if !ok {
+			break
+		}
+		hasJoinTable = true
+		from = append(from, t)
+		tableRefs = append(tableRefs, tableRef)
+		colRefs = append(colRefs, cols...)
 	}
 
 	update := &tree.Update{
 		Table:     table,
-		Where:     s.makeWhere(tableRefs, false /* hasJoinTable */),
-		OrderBy:   s.makeOrderBy(tableRefs),
+		From:      from,
+		Where:     s.makeWhere(colRefs, hasJoinTable),
+		OrderBy:   s.makeOrderBy(colRefs),
 		Limit:     makeLimit(s),
 		Returning: &tree.NoReturningClause{},
 	}
-	// Each row can be set at most once. Copy tableRefs to upRefs and remove
-	// elements from it as we use them.
-	upRefs := tableRefs.extend()
+	cols := make(map[tree.Name]*tree.ColumnTableDef)
+	for _, c := range ref.Columns {
+		cols[c.Name] = c
+	}
 	for (len(update.Exprs) < 1 || s.coin()) && len(upRefs) > 0 {
 		n := s.rnd.Intn(len(upRefs))
 		ref := upRefs[n]
@@ -950,7 +985,7 @@ func (s *Smither) makeUpdate(refs colRefs) (*tree.Update, *tableRef, bool) {
 		}
 		var expr tree.TypedExpr
 		for {
-			expr = makeScalar(s, ref.typ, tableRefs)
+			expr = makeScalar(s, ref.typ, colRefs)
 			// Make sure expr isn't null if that's not allowed.
 			if col.Nullable.Nullability != tree.NotNull || expr != tree.DNull {
 				break
@@ -968,7 +1003,7 @@ func (s *Smither) makeUpdate(refs colRefs) (*tree.Update, *tableRef, bool) {
 		update.OrderBy = nil
 	}
 
-	return update, tableRef, true
+	return update, tableRefs, true
 }
 
 func makeUpdateReturning(s *Smither, refs colRefs, forJoin bool) (tree.TableExpr, colRefs, bool) {
@@ -1157,7 +1192,7 @@ func (s *Smither) makeInsertReturning(refs colRefs) (tree.TableExpr, colRefs, bo
 		return nil, nil, false
 	}
 	var returningRefs colRefs
-	insert.Returning, returningRefs = s.makeReturning(insertRef)
+	insert.Returning, returningRefs = s.makeReturning([]*tableRef{insertRef})
 	return &tree.StatementSource{
 		Statement: insert,
 	}, returningRefs, true
@@ -1315,14 +1350,16 @@ func makeLimit(s *Smither) *tree.Limit {
 	return nil
 }
 
-func (s *Smither) makeReturning(table *tableRef) (*tree.ReturningExprs, colRefs) {
+func (s *Smither) makeReturning(tables []*tableRef) (*tree.ReturningExprs, colRefs) {
 	desiredTypes := s.makeDesiredTypes()
 
-	refs := make(colRefs, len(table.Columns))
-	for i, c := range table.Columns {
-		refs[i] = &colRef{
-			typ:  tree.MustBeStaticallyKnownType(c.Type),
-			item: &tree.ColumnItem{ColumnName: c.Name},
+	var refs colRefs
+	for _, table := range tables {
+		for _, c := range table.Columns {
+			refs = append(refs, &colRef{
+				typ:  tree.MustBeStaticallyKnownType(c.Type),
+				item: &tree.ColumnItem{ColumnName: c.Name},
+			})
 		}
 	}