tree: add pretty-printing for CreateRoutine

The routine body in `tree.CreateRoutine` is stored as a raw string and
written as-is by its `Format` method. We implement a new `Doc` method
that first re-parses the body so that it can be pretty-printed.

Release note (sql change): Statement bundles for CREATE FUNCTION and
CREATE PROCEDURE will now contain pretty-printed function bodies.

Author: andyyang890

pkg/sql/exec_util.go

@@ -159,6 +159,24 @@ func DoParserInjection() {
 	parserutils.ParseOne = parser.ParseOne
 	parserutils.ParseQualifiedTableName = parser.ParseQualifiedTableName
 	parserutils.PLpgSQLParse = plpgsqlparser.Parse
+	tree.ParseSQLForDoc = func(sql string) []tree.NodeFormatter {
+		stmts, err := parser.Parse(sql)
+		if err != nil {
+			return nil
+		}
+		result := make([]tree.NodeFormatter, len(stmts))
+		for i, s := range stmts {
+			result[i] = s.AST
+		}
+		return result
+	}
+	tree.ParsePLpgSQLForDoc = func(sql string) tree.NodeFormatter {
+		stmt, err := plpgsqlparser.Parse(sql)
+		if err != nil {
+			return nil
+		}
+		return stmt.AST
+	}
 }
 
 // ClusterOrganization is the organization name.

pkg/sql/sem/tree/create_routine.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
+	"github.com/cockroachdb/cockroach/pkg/util/pretty"
 	"github.com/cockroachdb/errors"
 )
 
@@ -80,6 +81,10 @@ type CreateRoutine struct {
 	Params      RoutineParams
 	ReturnType  *RoutineReturnType
 	Options     RoutineOptions
+	// RoutineBody is set during parsing for routines that use SQL-standard
+	// inline body syntax: either BEGIN ATOMIC ... END or a bare RETURN expr
+	// (as opposed to the dollar-quoted AS $$ ... $$ syntax). Currently
+	// unimplemented in the optimizer.
 	RoutineBody *RoutineBody
 	// BodyStatements is not assigned during initial parsing of user input. It's
 	// assigned during opt builder for logging purpose at the moment. It stores
@@ -165,6 +170,114 @@ func (node *CreateRoutine) Format(ctx *FmtCtx) {
 	}
 }
 
+// Doc implements the Docer interface.
+func (node *CreateRoutine) Doc(p *PrettyCfg) pretty.Doc {
+	header := pretty.Keyword("CREATE")
+	if node.Replace {
+		header = pretty.ConcatSpace(header, pretty.Keyword("OR REPLACE"))
+	}
+	if node.IsProcedure {
+		header = pretty.ConcatSpace(header, pretty.Keyword("PROCEDURE"))
+	} else {
+		header = pretty.ConcatSpace(header, pretty.Keyword("FUNCTION"))
+	}
+	header = pretty.ConcatSpace(header, p.Doc(&node.Name))
+	paramDocs := make([]pretty.Doc, len(node.Params))
+	for i := range node.Params {
+		paramDocs[i] = p.Doc(&node.Params[i])
+	}
+	header = pretty.Concat(header, p.bracket("(", p.commaSeparated(paramDocs...), ")"))
+
+	var clauses []pretty.Doc
+	if !node.IsProcedure && node.ReturnType != nil {
+		ret := pretty.Keyword("RETURNS")
+		if node.ReturnType.SetOf {
+			ret = pretty.ConcatSpace(ret, pretty.Keyword("SETOF"))
+		}
+		ret = pretty.ConcatSpace(ret, p.FormatType(node.ReturnType.Type))
+		clauses = append(clauses, ret)
+	}
+
+	// Extract RoutineBodyStr from options; it is formatted separately as the body.
+	var isPLpgSQL bool
+	var funcBody RoutineBodyStr
+	for _, option := range node.Options {
+		switch t := option.(type) {
+		case RoutineBodyStr:
+			funcBody = t
+			continue
+		case RoutineLeakproof, RoutineVolatility, RoutineNullInputBehavior:
+			if node.IsProcedure {
+				continue
+			}
+		case RoutineLanguage:
+			isPLpgSQL = t == RoutineLangPLpgSQL
+		}
+		clauses = append(clauses, p.Doc(option))
+	}
+
+	var bodyDoc pretty.Doc
+	if node.RoutineBody != nil {
+		stmtDocs := make([]pretty.Doc, len(node.RoutineBody.Stmts))
+		for i, stmt := range node.RoutineBody.Stmts {
+			d := p.Doc(stmt)
+			// DoBlock.Doc already includes a trailing semicolon.
+			if _, ok := stmt.(*DoBlock); !ok {
+				d = pretty.Concat(d, pretty.Text(";"))
+			}
+			stmtDocs[i] = d
+		}
+		bodyDoc = pretty.Fold(pretty.Concat,
+			pretty.Keyword("BEGIN ATOMIC"),
+			pretty.NestT(pretty.Concat(pretty.HardLine, pretty.Stack(stmtDocs...))),
+			pretty.HardLine,
+			pretty.Keyword("END"))
+	} else if !isPLpgSQL && ParseSQLForDoc != nil {
+		if stmts := ParseSQLForDoc(string(funcBody)); len(stmts) > 0 {
+			stmtDocs := make([]pretty.Doc, len(stmts))
+			for i, stmt := range stmts {
+				d := p.Doc(stmt)
+				// DoBlock.Doc already includes a trailing semicolon.
+				if _, ok := stmt.(*DoBlock); !ok {
+					d = pretty.Concat(d, pretty.Text(";"))
+				}
+				stmtDocs[i] = d
+			}
+			// Render the body to find a dollar-quote tag that doesn't
+			// collide with nested dollar quotes (e.g. from DO blocks).
+			fmtCtx := NewFmtCtx(p.FmtFlagsWithDefaults())
+			for _, stmt := range stmts {
+				fmtCtx.FormatNode(stmt)
+			}
+			bodyStr := fmtCtx.CloseAndGetString()
+			tag := "$" + DollarQuoteDelimiter(bodyStr) + "$"
+			bodyDoc = pretty.Fold(pretty.Concat,
+				pretty.ConcatSpace(pretty.Keyword("AS"), pretty.Text(tag)),
+				pretty.NestT(pretty.Concat(pretty.HardLine, pretty.Stack(stmtDocs...))),
+				pretty.HardLine,
+				pretty.Text(tag))
+		}
+	} else if isPLpgSQL && ParsePLpgSQLForDoc != nil {
+		if parsed := ParsePLpgSQLForDoc(string(funcBody)); parsed != nil {
+			// Use the formatted output to pick a dollar-quote tag that
+			// doesn't collide with nested dollar quotes (e.g. DO blocks).
+			bodyStr := AsStringWithFlags(parsed, p.FmtFlagsWithDefaults())
+			tag := "$" + DollarQuoteDelimiter(bodyStr) + "$"
+			bodyDoc = pretty.Fold(pretty.Concat,
+				pretty.ConcatSpace(pretty.Keyword("AS"), pretty.Text(tag)),
+				pretty.NestT(pretty.Concat(pretty.HardLine, p.Doc(parsed))),
+				pretty.HardLine,
+				pretty.Text(tag))
+		}
+	}
+	if bodyDoc == nil {
+		bodyDoc = p.docAsString(funcBody)
+	}
+	clauses = append(clauses, bodyDoc)
+
+	return p.nestUnder(header, pretty.Stack(clauses...))
+}
+
 // RoutineBody represent a list of statements in a UDF body.
 type RoutineBody struct {
 	// Stmts is populated during parsing. Unlike BodyStatements, we don't need

pkg/sql/sem/tree/pretty.go

@@ -20,6 +20,22 @@ import (
 	"github.com/cockroachdb/errors"
 )
 
+var (
+	// ParseSQLForDoc parses SQL statements stored as raw strings (e.g.
+	// RoutineBodyStr) into AST nodes for structured pretty-printing.
+	// Injected at init time to avoid a circular dependency on the parser
+	// package (the parserutils package unfortunately has the same issue).
+	// Returns nil if parsing fails.
+	ParseSQLForDoc func(sql string) []NodeFormatter
+
+	// ParsePLpgSQLForDoc parses a PL/pgSQL block stored as a raw string
+	// into an AST node for structured pretty-printing. Injected at init
+	// time to avoid a circular dependency on the plpgsqlparser package
+	// (the parserutils package unfortunately has the same issue).
+	// Returns nil if parsing fails.
+	ParsePLpgSQLForDoc func(sql string) NodeFormatter
+)
+
 // This file contains methods that convert statements to pretty Docs (a tree
 // structure that can be pretty printed at a specific line width). Nodes
 // implement the Docer interface to allow this conversion. In general,