Merge pull request #3017 from dolthub/angela/returning

Support `*` in `returning` clauses

Author: angelamayxie

enginetest/queries/insert_queries.go

@@ -2301,7 +2301,7 @@ var InsertScripts = []ScriptTest{
 				Expected: []sql.Row{{1, "Cat"}},
 			},
 			{
-				Query:    "insert into auto_pk values (NULL, 'Dog'),(5, 'Fish'),(NULL, 'Horse') returning pk,name",
+				Query:    "insert into auto_pk values (NULL, 'Dog'),(5, 'Fish'),(NULL, 'Horse') returning *",
 				Expected: []sql.Row{{2, "Dog"}, {5, "Fish"}, {6, "Horse"}},
 			},
 		},

sql/planbuilder/dml.go

@@ -150,12 +150,10 @@ func (b *Builder) buildInsert(inScope *scope, i *ast.Insert) (outScope *scope) {
 	ins := plan.NewInsertInto(db, plan.NewInsertDestination(sch, dest), srcScope.node, isReplace, columns, onDupExprs, ignore)
 	ins.LiteralValueSource = srcLiteralOnly
 
-	if i.Returning != nil {
-		returningExprs := make([]sql.Expression, len(i.Returning))
-		for i, selectExpr := range i.Returning {
-			returningExprs[i] = b.selectExprToExpression(destScope, selectExpr)
-		}
-		ins.Returning = returningExprs
+	if len(i.Returning) > 0 {
+		// TODO: read returning results from outScope instead of ins.Returning so that there is no need to return list
+		// of expressions
+		ins.Returning = b.analyzeSelectList(destScope, destScope, i.Returning)
 	}
 
 	b.validateInsert(ins)
@@ -583,11 +581,7 @@ func (b *Builder) buildUpdate(inScope *scope, u *ast.Update) (outScope *scope) {
 	}
 
 	if len(u.Returning) > 0 {
-		returningExprs := make([]sql.Expression, len(u.Returning))
-		for i, selectExpr := range u.Returning {
-			returningExprs[i] = b.selectExprToExpression(outScope, selectExpr)
-		}
-		update.Returning = returningExprs
+		update.Returning = b.analyzeSelectList(outScope, outScope, u.Returning)
 	}
 
 	outScope.node = update.WithChecks(checks)

sql/planbuilder/project.go

@@ -29,8 +29,8 @@ func (b *Builder) analyzeProjectionList(inScope, outScope *scope, selectExprs as
 	b.analyzeSelectList(inScope, outScope, selectExprs)
 }
 
-func (b *Builder) analyzeSelectList(inScope, outScope *scope, selectExprs ast.SelectExprs) {
-	// todo ideally we would not create new expressions here.
+func (b *Builder) analyzeSelectList(inScope, outScope *scope, selectExprs ast.SelectExprs) (expressions []sql.Expression) {
+	// TODO: ideally we would not create new expressions here.
 	// we want to in-place identify aggregations, expand stars.
 	// use inScope to construct projections for projScope
 
@@ -160,6 +160,7 @@ func (b *Builder) analyzeSelectList(inScope, outScope *scope, selectExprs ast.Se
 	}
 
 	inScope.parent = tempScope.parent
+	return exprs
 }
 
 // selectExprToExpression binds dependencies in a scalar expression in a SELECT clause.