release-23.2: sql: support sequence and udt name rewriting in plpgsql

pkg/sql/BUILD.bazel

@@ -490,6 +490,8 @@ go_library(
         "//pkg/sql/sem/catconstants",
         "//pkg/sql/sem/catid",
         "//pkg/sql/sem/eval",
+        "//pkg/sql/sem/plpgsqltree",
+        "//pkg/sql/sem/plpgsqltree/utils",
         "//pkg/sql/sem/semenumpb",
         "//pkg/sql/sem/transform",
         "//pkg/sql/sem/tree",

pkg/sql/crdb_internal.go

@@ -70,7 +70,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgnotice"
-	plpgsql "github.com/cockroachdb/cockroach/pkg/sql/plpgsql/parser"
 	"github.com/cockroachdb/cockroach/pkg/sql/privilege"
 	"github.com/cockroachdb/cockroach/pkg/sql/protoreflect"
 	"github.com/cockroachdb/cockroach/pkg/sql/roleoption"
@@ -3585,29 +3584,15 @@ func createRoutinePopulate(
 			for i := range treeNode.Options {
 				if body, ok := treeNode.Options[i].(tree.RoutineBodyStr); ok {
 					bodyStr := string(body)
-					switch fnDesc.GetLanguage() {
-					case catpb.Function_SQL:
-						bodyStr, err = formatFunctionQueryTypesForDisplay(ctx, &p.semaCtx, p.SessionData(), bodyStr)
-						if err != nil {
-							return err
-						}
-						bodyStr, err = formatQuerySequencesForDisplay(ctx, &p.semaCtx, bodyStr, true /* multiStmt */)
-						if err != nil {
-							return err
-						}
-						bodyStr = "\n" + bodyStr + "\n"
-					case catpb.Function_PLPGSQL:
-						// TODO(drewk): integrate this with the SQL case above.
-						plpgsqlStmt, err := plpgsql.Parse(bodyStr)
-						if err != nil {
-							return err
-						}
-						fmtCtx := tree.NewFmtCtx(tree.FmtParsable)
-						fmtCtx.FormatNode(plpgsqlStmt.AST)
-						bodyStr = "\n" + fmtCtx.CloseAndGetString()
-					default:
-						return errors.AssertionFailedf("unexpected function language: %s", fnDesc.GetLanguage())
+					bodyStr, err = formatFunctionQueryTypesForDisplay(ctx, &p.semaCtx, p.SessionData(), bodyStr, fnDesc.GetLanguage())
+					if err != nil {
+						return err
+					}
+					bodyStr, err = formatQuerySequencesForDisplay(ctx, &p.semaCtx, bodyStr, true /* multiStmt */, fnDesc.GetLanguage())
+					if err != nil {
+						return err
 					}
+					bodyStr = "\n" + bodyStr + "\n"
 					stmtStrs := strings.Split(bodyStr, "\n")
 					for i := range stmtStrs {
 						if stmtStrs[i] != "" {

pkg/sql/create_function.go

@@ -455,26 +455,19 @@ func setFuncOptions(
 		}
 	}
 
-	switch lang {
-	case catpb.Function_SQL:
+	if lang != catpb.Function_UNKNOWN_LANGUAGE && body != "" {
 		// Replace any sequence names in the function body with IDs.
-		seqReplacedFuncBody, err := replaceSeqNamesWithIDs(params.ctx, params.p, body, true)
+		seqReplacedFuncBody, err := replaceSeqNamesWithIDsLang(params.ctx, params.p, body, true, lang)
 		if err != nil {
 			return err
 		}
-		typeReplacedFuncBody, err := serializeUserDefinedTypes(
-			params.ctx, params.p.SemaCtx(), seqReplacedFuncBody, true /* multiStmt */, "UDFs",
-		)
+		typeReplacedFuncBody, err := serializeUserDefinedTypesLang(
+			params.ctx, params.p.SemaCtx(), seqReplacedFuncBody, true /* multiStmt */, "UDFs", lang)
 		if err != nil {
 			return err
 		}
 		udfDesc.SetFuncBody(typeReplacedFuncBody)
-	case catpb.Function_PLPGSQL:
-		// TODO(#115627): make replaceSeqNamesWithIDs and serializeUserDefinedTypes
-		// play nice with PL/pgSQL.
-		udfDesc.SetFuncBody(body)
 	}
-
 	return nil
 }
 

pkg/sql/create_view.go

@@ -30,8 +30,11 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgnotice"
+	plpgsql "github.com/cockroachdb/cockroach/pkg/sql/plpgsql/parser"
 	"github.com/cockroachdb/cockroach/pkg/sql/privilege"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/plpgsqltree"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/plpgsqltree/utils"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlerrors"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqltelemetry"
@@ -425,14 +428,31 @@ func makeViewTableDesc(
 	return desc, nil
 }
 
-// replaceSeqNamesWithIDs prepares to walk the given viewQuery by defining the
-// function used to replace sequence names with IDs, and parsing the
-// viewQuery into a statement.
+// replaceSeqNamesWithIDsLang walks the query in queryStr, replacing any
+// sequence names with their IDs and returning a new query string with the names
+// replaced. It assumes that the query is in the SQL language.
 // TODO (Chengxiong): move this to a better place.
 func replaceSeqNamesWithIDs(
 	ctx context.Context, sc resolver.SchemaResolver, queryStr string, multiStmt bool,
+) (string, error) {
+	return replaceSeqNamesWithIDsLang(ctx, sc, queryStr, multiStmt, catpb.Function_SQL)
+}
+
+// replaceSeqNamesWithIDsLang walks the query in queryStr, replacing any
+// sequence names with their IDs and returning a new query string with the names
+// replaced. Queries may be in either the SQL or PLpgSQL language, indicated by
+// lang.
+func replaceSeqNamesWithIDsLang(
+	ctx context.Context,
+	sc resolver.SchemaResolver,
+	queryStr string,
+	multiStmt bool,
+	lang catpb.Function_Language,
 ) (string, error) {
 	replaceSeqFunc := func(expr tree.Expr) (recurse bool, newExpr tree.Expr, err error) {
+		if expr == nil {
+			return false, expr, nil
+		}
 		seqIdentifiers, err := seqexpr.GetUsedSequences(expr)
 		if err != nil {
 			return false, expr, err
@@ -452,48 +472,83 @@ func replaceSeqNamesWithIDs(
 		return false, newExpr, nil
 	}
 
-	var stmts tree.Statements
-	if multiStmt {
-		parsedStmtd, err := parser.Parse(queryStr)
-		if err != nil {
-			return "", errors.Wrap(err, "failed to parse query string")
+	fmtCtx := tree.NewFmtCtx(tree.FmtSimple)
+	switch lang {
+	case catpb.Function_SQL:
+		var stmts tree.Statements
+		if multiStmt {
+			parsedStmtd, err := parser.Parse(queryStr)
+			if err != nil {
+				return "", errors.Wrap(err, "failed to parse query string")
+			}
+			for _, s := range parsedStmtd {
+				stmts = append(stmts, s.AST)
+			}
+		} else {
+			stmt, err := parser.ParseOne(queryStr)
+			if err != nil {
+				return "", errors.Wrap(err, "failed to parse query string")
+			}
+			stmts = tree.Statements{stmt.AST}
 		}
-		for _, s := range parsedStmtd {
-			stmts = append(stmts, s.AST)
+
+		for i, stmt := range stmts {
+			newStmt, err := tree.SimpleStmtVisit(stmt, replaceSeqFunc)
+			if err != nil {
+				return "", err
+			}
+			if i > 0 {
+				fmtCtx.WriteString("\n")
+			}
+			fmtCtx.FormatNode(newStmt)
+			if multiStmt {
+				fmtCtx.WriteString(";")
+			}
 		}
-	} else {
-		stmt, err := parser.ParseOne(queryStr)
+	case catpb.Function_PLPGSQL:
+		var stmts plpgsqltree.Statement
+		plstmt, err := plpgsql.Parse(queryStr)
 		if err != nil {
 			return "", errors.Wrap(err, "failed to parse query string")
 		}
-		stmts = tree.Statements{stmt.AST}
-	}
+		stmts = plstmt.AST
 
-	fmtCtx := tree.NewFmtCtx(tree.FmtSimple)
-	for i, stmt := range stmts {
-		newStmt, err := tree.SimpleStmtVisit(stmt, replaceSeqFunc)
-		if err != nil {
-			return "", err
-		}
-		if i > 0 {
-			fmtCtx.WriteString("\n")
-		}
+		v := utils.SQLStmtVisitor{Fn: replaceSeqFunc}
+		newStmt := plpgsqltree.Walk(&v, stmts)
 		fmtCtx.FormatNode(newStmt)
-		if multiStmt {
-			fmtCtx.WriteString(";")
-		}
 	}
 
 	return fmtCtx.String(), nil
 }
 
-// serializeUserDefinedTypes will walk the given view query
-// and serialize any user defined types, so that renaming the type
-// does not corrupt the view.
+// serializeUserDefinedTypes walks the given query and serializes any
+// user defined types as IDs, so that renaming the type does not cause
+// corruption, and returns a new query string containing the replacement IDs.
+// It assumes that the query language is SQL.
 func serializeUserDefinedTypes(
 	ctx context.Context, semaCtx *tree.SemaContext, queries string, multiStmt bool, parentType string,
 ) (string, error) {
+	return serializeUserDefinedTypesLang(ctx, semaCtx, queries, multiStmt, parentType, catpb.Function_SQL)
+}
+
+// serializeUserDefinedTypesLang walks the given query and serializes any
+// user defined types as IDs, so that renaming the type does not cause
+// corruption, and returns a new query string containing the replacement IDs.
+// The query may be in either the SQL or PLpgSQL language, indicated by lang.
+func serializeUserDefinedTypesLang(
+	ctx context.Context,
+	semaCtx *tree.SemaContext,
+	queries string,
+	multiStmt bool,
+	parentType string,
+	lang catpb.Function_Language,
+) (string, error) {
+	// replaceFunc is a visitor function that replaces user defined types in SQL
+	// expressions with their IDs.
 	replaceFunc := func(expr tree.Expr) (recurse bool, newExpr tree.Expr, err error) {
+		if expr == nil {
+			return false, expr, nil
+		}
 		var innerExpr tree.Expr
 		var typRef tree.ResolvableTypeReference
 		switch n := expr.(type) {
@@ -539,39 +594,83 @@ func serializeUserDefinedTypes(
 		}
 		return false, parsedExpr, nil
 	}
-
-	var stmts tree.Statements
-	if multiStmt {
-		parsedStmts, err := parser.Parse(queries)
-		if err != nil {
-			return "", errors.Wrap(err, "failed to parse query")
+	// replaceTypeFunc is a visitor function that replaces type annotations
+	// containing user defined types with their IDs. This is currently only
+	// necessary for some kinds of PLpgSQL statements.
+	replaceTypeFunc := func(typ tree.ResolvableTypeReference) (newTyp tree.ResolvableTypeReference, err error) {
+		if typ == nil {
+			return typ, nil
 		}
-		stmts = make(tree.Statements, len(parsedStmts))
-		for i, stmt := range parsedStmts {
-			stmts[i] = stmt.AST
+		// semaCtx may be nil if this is a virtual view being created at
+		// init time.
+		var typeResolver tree.TypeReferenceResolver
+		if semaCtx != nil {
+			typeResolver = semaCtx.TypeResolver
 		}
-	} else {
-		stmt, err := parser.ParseOne(queries)
+		var t *types.T
+		t, err = tree.ResolveType(ctx, typ, typeResolver)
 		if err != nil {
-			return "", errors.Wrap(err, "failed to parse query")
+			return typ, err
+		}
+		if !t.UserDefined() {
+			return typ, nil
 		}
-		stmts = tree.Statements{stmt.AST}
+		return &tree.OIDTypeReference{OID: t.Oid()}, nil
 	}
 
 	fmtCtx := tree.NewFmtCtx(tree.FmtSimple)
-	for i, stmt := range stmts {
-		newStmt, err := tree.SimpleStmtVisit(stmt, replaceFunc)
-		if err != nil {
-			return "", err
+	switch lang {
+	case catpb.Function_SQL:
+		var stmts tree.Statements
+		if multiStmt {
+			parsedStmts, err := parser.Parse(queries)
+			if err != nil {
+				return "", errors.Wrap(err, "failed to parse query")
+			}
+			stmts = make(tree.Statements, len(parsedStmts))
+			for i, stmt := range parsedStmts {
+				stmts[i] = stmt.AST
+			}
+		} else {
+			stmt, err := parser.ParseOne(queries)
+			if err != nil {
+				return "", errors.Wrap(err, "failed to parse query")
+			}
+			stmts = tree.Statements{stmt.AST}
 		}
-		if i > 0 {
-			fmtCtx.WriteString("\n")
+
+		for i, stmt := range stmts {
+			newStmt, err := tree.SimpleStmtVisit(stmt, replaceFunc)
+			if err != nil {
+				return "", err
+			}
+			if i > 0 {
+				fmtCtx.WriteString("\n")
+			}
+			fmtCtx.FormatNode(newStmt)
+			if multiStmt {
+				fmtCtx.WriteString(";")
+			}
 		}
-		fmtCtx.FormatNode(newStmt)
-		if multiStmt {
-			fmtCtx.WriteString(";")
+	case catpb.Function_PLPGSQL:
+		var stmts plpgsqltree.Statement
+		plstmt, err := plpgsql.Parse(queries)
+		if err != nil {
+			return "", errors.Wrap(err, "failed to parse query string")
 		}
+		stmts = plstmt.AST
+
+		v := utils.SQLStmtVisitor{Fn: replaceFunc}
+		newStmt := plpgsqltree.Walk(&v, stmts)
+		// Some PLpgSQL statements (i.e., declarations), may contain type
+		// annotations containing the UDT. We need to walk the AST to replace them,
+		// too.
+		v2 := utils.TypeRefVisitor{Fn: replaceTypeFunc}
+		newStmt = plpgsqltree.Walk(&v2, newStmt)
+		fmtCtx.FormatNode(newStmt)
+		fmtCtx.WriteString(";")
 	}
+
 	return fmtCtx.CloseAndGetString(), nil
 }
 

pkg/sql/opt/optbuilder/plpgsql.go

@@ -1531,7 +1531,7 @@ func newRecordTypeVisitor(
 
 var _ ast.StatementVisitor = &recordTypeVisitor{}
 
-func (r *recordTypeVisitor) Visit(stmt ast.Statement) {
+func (r *recordTypeVisitor) Visit(stmt ast.Statement) (newStmt ast.Statement, changed bool) {
 	if retStmt, ok := stmt.(*ast.Return); ok {
 		desired := types.Any
 		if r.typ != types.Unknown {
@@ -1544,7 +1544,7 @@ func (r *recordTypeVisitor) Visit(stmt ast.Statement) {
 		}
 		typ := typedExpr.ResolvedType()
 		if typ == types.Unknown {
-			return
+			return stmt, false
 		}
 		if typ.Family() != types.TupleFamily {
 			panic(pgerror.New(pgcode.DatatypeMismatch,
@@ -1553,7 +1553,7 @@ func (r *recordTypeVisitor) Visit(stmt ast.Statement) {
 		}
 		if r.typ == types.Unknown {
 			r.typ = typ
-			return
+			return stmt, false
 		}
 		if !typ.Identical(r.typ) {
 			panic(errors.WithHint(
@@ -1564,4 +1564,5 @@ func (r *recordTypeVisitor) Visit(stmt ast.Statement) {
 			))
 		}
 	}
+	return stmt, false
 }

pkg/sql/plpgsql/parser/testdata/stmt_case

@@ -133,6 +133,7 @@ BEGIN
 END CASE;
 END
 ----
+decl_stmt: 1
 stmt_block: 1
 stmt_call: 3
 stmt_case: 1

pkg/sql/sem/plpgsqltree/BUILD.bazel

@@ -13,6 +13,7 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/sql/sem/tree",
+        "//pkg/util/errorutil/unimplemented",
         "@com_github_cockroachdb_errors//:errors",
     ],
 )