returning-star (#556)

Summary:

- Support for `RETURNING *`.
- Added robot test `Insert Returning Star`.
- Added robot test `Insert Async Returning Star`.
- Added robot test `Update Returning Star`.
- Added robot test `Update Async Returning Star`.
- Added robot test `Replace Returning Star`.
- Added robot test `Replace Async Returning Star`.

Author: general-kroll-4-life

go.mod

@@ -19,7 +19,7 @@ require (
 	github.com/spf13/cobra v1.4.0
 	github.com/spf13/pflag v1.0.5
 	github.com/spf13/viper v1.10.1
-	github.com/stackql/any-sdk v0.1.4-alpha12
+	github.com/stackql/any-sdk v0.1.4-beta03
 	github.com/stackql/go-suffix-map v0.0.1-alpha01
 	github.com/stackql/psql-wire v0.1.1-beta23
 	github.com/stackql/stackql-parser v0.0.15-alpha06

go.sum

@@ -484,8 +484,8 @@ github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/spf13/viper v1.10.1 h1:nuJZuYpG7gTj/XqiUwg8bA0cp1+M2mC3J4g5luUYBKk=
 github.com/spf13/viper v1.10.1/go.mod h1:IGlFPqhNAPKRxohIzWpI5QEy4kuI7tcl5WvR+8qy1rU=
-github.com/stackql/any-sdk v0.1.4-alpha12 h1:SozoNb4+7+y+0bVGmO8ozFtTTpYhEWBHVDuN9SmkFuA=
-github.com/stackql/any-sdk v0.1.4-alpha12/go.mod h1:AKS/g28y7m4SWL/YW8veE9MCNy8XJgaicVibemVE9e8=
+github.com/stackql/any-sdk v0.1.4-beta03 h1:hEpVc04jQ2BgfiyIxBVB1K4cyzjlc2BL6dH/Of4jRMU=
+github.com/stackql/any-sdk v0.1.4-beta03/go.mod h1:AKS/g28y7m4SWL/YW8veE9MCNy8XJgaicVibemVE9e8=
 github.com/stackql/go-suffix-map v0.0.1-alpha01 h1:TDUDS8bySu41Oo9p0eniUeCm43mnRM6zFEd6j6VUaz8=
 github.com/stackql/go-suffix-map v0.0.1-alpha01/go.mod h1:QAi+SKukOyf4dBtWy8UMy+hsXXV+yyEE4vmBkji2V7g=
 github.com/stackql/psql-wire v0.1.1-beta23 h1:1ayYMjZArfDcIMyEOKnm+Bp1zRCISw8pguvTFuUhhVQ=

internal/stackql/parserutil/parser_util.go

@@ -90,11 +90,24 @@ func ExtractSelectColumnNames(selStmt *sqlparser.Select, formatter sqlparser.Nod
 
 func ExtractInsertReturningColumnNames(
 	insertStmt *sqlparser.Insert,
+	starColumns []string,
+	formatter sqlparser.NodeFormatter,
+) ([]ColumnHandle, error) {
+	return extractInsertReturningColumnNames(
+		insertStmt.SelectExprs,
+		starColumns,
+		formatter,
+	)
+}
+
+func extractInsertReturningColumnNames(
+	returningExprs sqlparser.SelectExprs,
+	starColumns []string,
 	formatter sqlparser.NodeFormatter,
 ) ([]ColumnHandle, error) {
 	var colNames []ColumnHandle
 	var err error
-	for _, node := range insertStmt.SelectExprs {
+	for _, node := range returningExprs {
 		switch node := node.(type) {
 		case *sqlparser.AliasedExpr:
 			cn, cErr := inferColNameFromExpr(node.Expr, formatter, node.As.GetRawVal())
@@ -103,29 +116,37 @@ func ExtractInsertReturningColumnNames(
 			}
 			colNames = append(colNames, cn)
 		case *sqlparser.StarExpr:
+			if len(starColumns) == 0 {
+				return nil, fmt.Errorf("no star columns provided for insert returning of star")
+			}
+			prefix := node.TableName.GetRawVal()
+			for _, colName := range starColumns {
+				if prefix != "" {
+					colName = prefix + "." + colName
+				}
+				colNames = append(colNames, ColumnHandle{
+					Name:  colName,
+					Alias: colName,
+				})
+			}
+		default:
+			err = fmt.Errorf("cannot use SelectExpr of type '%T' as a raw value", node)
+			return nil, err
 		}
 	}
 	return colNames, err
 }
 
 func ExtractUpdateReturningColumnNames(
 	updateStmt *sqlparser.Update,
+	starColumns []string,
 	formatter sqlparser.NodeFormatter,
 ) ([]ColumnHandle, error) {
-	var colNames []ColumnHandle
-	var err error
-	for _, node := range updateStmt.SelectExprs {
-		switch node := node.(type) {
-		case *sqlparser.AliasedExpr:
-			cn, cErr := inferColNameFromExpr(node.Expr, formatter, node.As.GetRawVal())
-			if cErr != nil {
-				return nil, cErr
-			}
-			colNames = append(colNames, cn)
-		case *sqlparser.StarExpr:
-		}
-	}
-	return colNames, err
+	return extractInsertReturningColumnNames(
+		updateStmt.SelectExprs,
+		starColumns,
+		formatter,
+	)
 }
 
 func ExtractInsertColumnNames(insertStmt *sqlparser.Insert) ([]string, error) {

internal/stackql/primitivegenerator/statement_analyzer.go

@@ -1152,7 +1152,8 @@ func (pb *standardPrimitiveGenerator) AnalyzeInsert(pbi planbuilderinput.PlanBui
 	}
 	columnHandles := []parserutil.ColumnHandle{}
 	if len(node.SelectExprs) > 0 {
-		columnHandles, err = parserutil.ExtractInsertReturningColumnNames(node, handlerCtx.GetASTFormatter())
+		starColumns, _ := methodAnalysisOutput.GetOrderedStarColumnsNames()
+		columnHandles, err = parserutil.ExtractInsertReturningColumnNames(node, starColumns, handlerCtx.GetASTFormatter())
 		if err != nil {
 			return err
 		}
@@ -1261,7 +1262,12 @@ func (pb *standardPrimitiveGenerator) AnalyzeUpdate(pbi planbuilderinput.PlanBui
 	}
 	columnHandles := []parserutil.ColumnHandle{}
 	if len(node.SelectExprs) > 0 {
-		columnHandles, err = parserutil.ExtractUpdateReturningColumnNames(node, handlerCtx.GetASTFormatter())
+		starColumns, _ := methodAnalysisOutput.GetOrderedStarColumnsNames()
+		columnHandles, err = parserutil.ExtractUpdateReturningColumnNames(
+			node,
+			starColumns,
+			handlerCtx.GetASTFormatter(),
+		)
 		if err != nil {
 			return err
 		}