sql: support for RETURNING *

This change adds support for RETURNING * in inserts and deletes. For updates we
currently only expand the updated columns; this is related to an existing
problem (we can't return other columns even explicitly; filed cockroachdb#4645).

Fixes cockroachdb#4593.

sql/logic_test.go

@@ -548,7 +548,12 @@ func (t *logicTest) execQuery(query logicQuery) {
 
 	var results []string
 	if query.colNames {
-		results = append(results, cols...)
+		for _, col := range cols {
+			// We split string results on whitespace and append a separate result
+			// for each string. A bit unusual, but otherwise we can't match strings
+			// containing whitespace.
+			results = append(results, strings.Fields(col)...)
+		}
 	}
 	for rows.Next() {
 		if err := rows.Scan(vals...); err != nil {

sql/returning.go

@@ -42,31 +42,38 @@ func newReturningHelper(p *planner, r parser.ReturningExprs,
 		return rh, nil
 	}
 
-	rh.results.columns = make([]ResultColumn, len(r))
+	rh.results.columns = make([]ResultColumn, 0, len(r))
 	table := tableInfo{
 		columns: makeResultColumns(tablecols, 0),
 		alias:   alias,
 	}
 	rh.qvals = make(qvalMap)
-	rh.exprs = make([]parser.Expr, len(r))
-	for i, c := range r {
-		expr, err := resolveQNames(&table, rh.qvals, c.Expr)
+	rh.exprs = make([]parser.Expr, 0, len(r))
+	for _, target := range r {
+
+		if isStar, cols, exprs, err := checkRenderStar(target, &table, rh.qvals); err != nil {
+			return returningHelper{}, err
+		} else if isStar {
+			rh.exprs = append(rh.exprs, exprs...)
+			rh.results.columns = append(rh.results.columns, cols...)
+			continue
+		}
+
+		// When generating an output column name it should exactly match the original
+		// expression, so determine the output column name before we perform any
+		// manipulations to the expression.
+		outputName := getRenderColName(target)
+
+		expr, err := resolveQNames(&table, rh.qvals, target.Expr)
 		if err != nil {
-			return rh, err
+			return returningHelper{}, err
 		}
-		rh.exprs[i] = expr
 		typ, err := expr.TypeCheck(rh.p.evalCtx.Args)
 		if err != nil {
-			return rh, err
-		}
-		name := string(c.As)
-		if name == "" {
-			name = expr.String()
-		}
-		rh.results.columns[i] = ResultColumn{
-			Name: name,
-			Typ:  typ,
+			return returningHelper{}, err
 		}
+		rh.exprs = append(rh.exprs, expr)
+		rh.results.columns = append(rh.results.columns, ResultColumn{Name: outputName, Typ: typ})
 	}
 	return rh, nil
 }

sql/select.go

@@ -370,51 +370,78 @@ func (s *selectNode) colIndex(expr parser.Expr) (int, error) {
 	}
 }
 
-func (s *selectNode) addRender(target parser.SelectExpr) *roachpb.Error {
-	// outputName will be empty if the target is not aliased.
-	outputName := string(target.As)
+// Checks if the SelectExpr is a QualifiedName with a StarIndirection prefix. If so, we match the
+// prefix of the qualified name to one of the tables in the query and then expand the "*" into a
+// list of columns. The qvalMap is updated to include all the relevant columns. A ResultColumns
+// and Expr pair is returned for each column.
+func checkRenderStar(target parser.SelectExpr, table *tableInfo, qvals qvalMap) (isStar bool,
+	columns []ResultColumn, exprs []parser.Expr, err error) {
 
-	switch t := target.Expr.(type) {
-	case *parser.QualifiedName:
-		// If a QualifiedName has a StarIndirection suffix we need to match the
-		// prefix of the qualified name to one of the tables in the query and
-		// then expand the "*" into a list of columns.
-		if s.pErr = roachpb.NewError(t.NormalizeColumnName()); s.pErr != nil {
-			return s.pErr
-		}
-		if t.IsStar() {
-			if s.table.alias == "" {
-				return roachpb.NewUErrorf("\"%s\" with no tables specified is not valid", t)
-			}
-			if target.As != "" {
-				return roachpb.NewUErrorf("\"%s\" cannot be aliased", t)
-			}
-			// TODO(radu): support multiple FROMs, consolidate with logic in findColumn
-			tableName := t.Table()
-			if tableName != "" && !equalName(s.table.alias, tableName) {
-				return roachpb.NewUErrorf("table \"%s\" not found", tableName)
-			}
+	qname, ok := target.Expr.(*parser.QualifiedName)
+	if !ok {
+		return false, nil, nil, nil
+	}
+	if err := qname.NormalizeColumnName(); err != nil {
+		return false, nil, nil, err
+	}
+	if !qname.IsStar() {
+		return false, nil, nil, nil
+	}
 
-			for idx, col := range s.table.columns {
-				if col.hidden {
-					continue
-				}
-				qval := s.qvals.getQVal(columnRef{&s.table, idx})
-				s.columns = append(s.columns, ResultColumn{Name: col.Name, Typ: qval.datum})
-				s.render = append(s.render, qval)
-			}
+	if table.alias == "" {
+		return false, nil, nil, fmt.Errorf("\"%s\" with no tables specified is not valid", qname)
+	}
+	if target.As != "" {
+		return false, nil, nil, fmt.Errorf("\"%s\" cannot be aliased", qname)
+	}
 
-			return nil
-		}
-	default:
-		if outputName == "" {
-			// When generating an output column name it should exactly match the original
-			// expression, so determine the output column name before we perform any
-			// manipulations to the expression (such as star expansion).
-			outputName = target.Expr.String()
+	// TODO(radu): support multiple FROMs, consolidate with logic in findColumn
+	tableName := qname.Table()
+	if tableName != "" && !equalName(table.alias, tableName) {
+		return false, nil, nil, fmt.Errorf("table \"%s\" not found", tableName)
+	}
+
+	for idx, col := range table.columns {
+		if col.hidden {
+			continue
 		}
+		qval := qvals.getQVal(columnRef{table, idx})
+		columns = append(columns, ResultColumn{Name: col.Name, Typ: qval.datum})
+		exprs = append(exprs, qval)
+	}
+	return true, columns, exprs, nil
+}
+
+// getRenderColName returns the output column name for a render expression.
+// The expression cannot be a star.
+func getRenderColName(target parser.SelectExpr) string {
+	if target.As != "" {
+		return string(target.As)
+	}
+	if qname, ok := target.Expr.(*parser.QualifiedName); ok {
+		return qname.Column()
+	}
+	return target.Expr.String()
+}
+
+func (s *selectNode) addRender(target parser.SelectExpr) *roachpb.Error {
+	// outputName will be empty if the target is not aliased.
+	outputName := string(target.As)
+
+	if isStar, cols, exprs, err := checkRenderStar(target, &s.table, s.qvals); err != nil {
+		s.pErr = roachpb.NewError(err)
+		return s.pErr
+	} else if isStar {
+		s.columns = append(s.columns, cols...)
+		s.render = append(s.render, exprs...)
+		return nil
 	}
 
+	// When generating an output column name it should exactly match the original
+	// expression, so determine the output column name before we perform any
+	// manipulations to the expression.
+	outputName = getRenderColName(target)
+
 	// Resolve qualified names. This has the side-effect of normalizing any
 	// qualified name found.
 	var resolved parser.Expr

sql/testdata/delete

@@ -78,3 +78,20 @@ DELETE FROM unindexed RETURNING k, v
 query II
 SELECT * FROM unindexed
 ----
+
+statement ok
+INSERT INTO unindexed VALUES (1, 2), (3, 4), (5, 6), (7, 8)
+
+query II colnames
+DELETE FROM unindexed WHERE k=3 or v=6 RETURNING *
+----
+k v
+3 4
+5 6
+
+query II colnames
+DELETE FROM unindexed RETURNING unindexed.*
+----
+k v
+1 2
+7 8

sql/testdata/insert

@@ -339,6 +339,38 @@ INSERT INTO return VALUES (default) RETURNING rowid != unique_rowid()
 statement error pq: qualified name "b" not found
 INSERT INTO return (a) VALUES (default) RETURNING b
 
-# TODO(mjibson): support *
-statement error pq: qualified name "\*" not found
+query II colnames
 INSERT INTO return VALUES (default) RETURNING *
+----
+a b
+3 NULL
+
+query II colnames
+INSERT INTO return VALUES (1) RETURNING *
+----
+a b
+1 NULL
+
+query II colnames
+INSERT INTO return VALUES (1, 2), (3, 4) RETURNING return.a, b
+----
+a b
+1 2
+3 4
+
+query II colnames
+INSERT INTO return VALUES (1, 2), (3, 4) RETURNING *
+----
+a b
+1 2
+3 4
+
+statement error pq: "return.*" cannot be aliased
+INSERT INTO return VALUES (1, 2), (3, 4) RETURNING return.* as x
+
+query III colnames
+INSERT INTO return VALUES (1, 2), (3, 4) RETURNING return.*, a + b
+----
+a b a + b
+1 2 3
+3 4 7

sql/testdata/update

@@ -171,11 +171,39 @@ SELECT * FROM abc
 ----
 1 6 7
 
-query III
-UPDATE abc SET (b, c) = (VALUES (8, 9)) RETURNING b, c, 4
+query III colnames
+UPDATE abc SET (b, c) = (VALUES (8, 9)) RETURNING abc.b, c, 4
 ----
+b c 4
 8 9 4
 
+query III colnames
+UPDATE abc SET (b, c) = (VALUES (8, 9)) RETURNING b as col1, c as col2, 4 as col3
+----
+col1 col2 col3
+8    9    4
+
+# Issue #4645: we should be able to return other columns.
+query error pq: qualified name "a" not found
+UPDATE abc SET (b, c) = (VALUES (8, 9)) RETURNING a
+
+# Issue #4645:  this should return all columns in the table.
+query II colnames
+UPDATE abc SET (b, c) = (VALUES (8, 9)) RETURNING *
+----
+b c
+8 9
+
+# Issue #4645:  this should return all columns in the table.
+query II colnames
+UPDATE abc SET (b, c) = (VALUES (8, 9)) RETURNING abc.*
+----
+b c
+8 9
+
+statement error pq: "abc.*" cannot be aliased
+UPDATE abc SET (b, c) = (VALUES (8, 9)) RETURNING abc.* as x
+
 query III
 SELECT * FROM abc
 ----