Add support for UPDATE ... RETURNING

go.mod

@@ -10,7 +10,7 @@ require (
 	github.com/dolthub/dolt/go/gen/proto/dolt/services/eventsapi v0.0.0-20241119094239-f4e529af734d
 	github.com/dolthub/flatbuffers/v23 v23.3.3-dh.2
 	github.com/dolthub/go-icu-regex v0.0.0-20250327004329-6799764f2dad
-	github.com/dolthub/go-mysql-server v0.19.1-0.20250423225005-2d6449f73e2b
+	github.com/dolthub/go-mysql-server v0.19.1-0.20250425170657-81264bff934b
 	github.com/dolthub/sqllogictest/go v0.0.0-20240618184124-ca47f9354216
 	github.com/dolthub/vitess v0.0.0-20250423221552-f731ee5c5379
 	github.com/fatih/color v1.13.0

go.sum

@@ -266,8 +266,8 @@ github.com/dolthub/fslock v0.0.3 h1:iLMpUIvJKMKm92+N1fmHVdxJP5NdyDK5bK7z7Ba2s2U=
 github.com/dolthub/fslock v0.0.3/go.mod h1:QWql+P17oAAMLnL4HGB5tiovtDuAjdDTPbuqx7bYfa0=
 github.com/dolthub/go-icu-regex v0.0.0-20250327004329-6799764f2dad h1:66ZPawHszNu37VPQckdhX1BPPVzREsGgNxQeefnlm3g=
 github.com/dolthub/go-icu-regex v0.0.0-20250327004329-6799764f2dad/go.mod h1:ylU4XjUpsMcvl/BKeRRMXSH7e7WBrPXdSLvnRJYrxEA=
-github.com/dolthub/go-mysql-server v0.19.1-0.20250423225005-2d6449f73e2b h1:kgKEFYkLTErT+VE4jqZsoEeq1hSKeXAOnPZA/8tDIyI=
-github.com/dolthub/go-mysql-server v0.19.1-0.20250423225005-2d6449f73e2b/go.mod h1:skB7lRdqgWVqDZ/ih/DtX8iSvK15IaxRf2ASliUwgtE=
+github.com/dolthub/go-mysql-server v0.19.1-0.20250425170657-81264bff934b h1:XtYHOUT9oDnyrT4n8N+3fdJnt8FidT8zAuOIpSHg4G0=
+github.com/dolthub/go-mysql-server v0.19.1-0.20250425170657-81264bff934b/go.mod h1:skB7lRdqgWVqDZ/ih/DtX8iSvK15IaxRf2ASliUwgtE=
 github.com/dolthub/gozstd v0.0.0-20240423170813-23a2903bca63 h1:OAsXLAPL4du6tfbBgK0xXHZkOlos63RdKYS3Sgw/dfI=
 github.com/dolthub/gozstd v0.0.0-20240423170813-23a2903bca63/go.mod h1:lV7lUeuDhH5thVGDCKXbatwKy2KW80L4rMT46n+Y2/Q=
 github.com/dolthub/ishell v0.0.0-20240701202509-2b217167d718 h1:lT7hE5k+0nkBdj/1UOSFwjWpNxf+LCApbRHgnCA17XE=

server/ast/update.go

@@ -24,16 +24,21 @@ import (
 )
 
 // nodeUpdate handles *tree.Update nodes.
-func nodeUpdate(ctx *Context, node *tree.Update) (*vitess.Update, error) {
+func nodeUpdate(ctx *Context, node *tree.Update) (update *vitess.Update, err error) {
 	if node == nil {
 		return nil, nil
 	}
 	ctx.Auth().PushAuthType(auth.AuthType_UPDATE)
 	defer ctx.Auth().PopAuthType()
 
-	if _, ok := node.Returning.(*tree.NoReturningClause); !ok {
-		return nil, errors.Errorf("RETURNING is not yet supported")
+	var returningExprs vitess.SelectExprs
+	if returning, ok := node.Returning.(*tree.ReturningExprs); ok {
+		returningExprs, err = nodeSelectExprs(ctx, tree.SelectExprs(*returning))
+		if err != nil {
+			return nil, err
+		}
 	}
+
 	if len(node.From) > 0 {
 		return nil, errors.Errorf("FROM is not yet supported")
 	}
@@ -68,5 +73,6 @@ func nodeUpdate(ctx *Context, node *tree.Update) (*vitess.Update, error) {
 		Where:      where,
 		OrderBy:    orderBy,
 		Limit:      limit,
+		Returning:  returningExprs,
 	}, nil
 }

server/connection_handler.go

@@ -1137,21 +1137,31 @@ func returnsRow(query ConvertedQuery) bool {
 	switch query.StatementTag {
 	case "SELECT", "SHOW", "FETCH", "EXPLAIN", "SHOW TABLES":
 		return true
-	case "INSERT":
-		hasReturningClause := false
-		sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) {
-			switch node := node.(type) {
-			case *sqlparser.Insert:
-				if len(node.Returning) > 0 {
-					hasReturningClause = true
-				}
-				return false, nil
-			}
-			// this should be impossible, but just in case
-			return true, nil
-		}, query.AST)
-		return hasReturningClause
+	case "INSERT", "UPDATE", "DELETE":
+		return hasReturningClause(query.AST)
 	default:
 		return false
 	}
 }
+
+// hasReturningClause return true if |statement| has a RETURNING clause defined.
+func hasReturningClause(statement sqlparser.Statement) bool {
+	hasReturningClause := false
+	sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) {
+		switch node := node.(type) {
+		case *sqlparser.Update:
+			if len(node.Returning) > 0 {
+				hasReturningClause = true
+			}
+			return false, nil
+		case *sqlparser.Insert:
+			if len(node.Returning) > 0 {
+				hasReturningClause = true
+			}
+			return false, nil
+		}
+		return true, nil
+	}, statement)
+
+	return hasReturningClause
+}