diff --git a/pkg/sql/logictest/testdata/logic_test/udf b/pkg/sql/logictest/testdata/logic_test/udf index 882709a2dec1..cb19306bb62e 100644 --- a/pkg/sql/logictest/testdata/logic_test/udf +++ b/pkg/sql/logictest/testdata/logic_test/udf @@ -3566,3 +3566,19 @@ query T SELECT f104927() ---- [{"i": 1, "s": "foo"}] + +subtest 104242 + +statement ok +CREATE TABLE t_104242 (i INT PRIMARY KEY); + +statement ok +CREATE TYPE typ_104242 AS ENUM ('a', 'b', 'c'); + +# We've only fixed issue with declarative schema changer. +skipif config local-legacy-schema-changer +skipif config local-mixed-22.2-23.1 +statement ok +CREATE FUNCTION f_104242() RETURNS INT LANGUAGE SQL AS $$ SELECT 1 FROM t_104242 WHERE NULL::typ_104242 IN () $$; + +subtest end diff --git a/pkg/sql/schemachanger/scbuild/BUILD.bazel b/pkg/sql/schemachanger/scbuild/BUILD.bazel index fb122d121d30..0cdb7654daea 100644 --- a/pkg/sql/schemachanger/scbuild/BUILD.bazel +++ b/pkg/sql/schemachanger/scbuild/BUILD.bazel @@ -33,6 +33,7 @@ go_library( "//pkg/sql/catalog/typedesc", "//pkg/sql/faketreeeval", "//pkg/sql/parser", + "//pkg/sql/parser/statements", "//pkg/sql/pgwire/pgcode", "//pkg/sql/pgwire/pgerror", "//pkg/sql/privilege", diff --git a/pkg/sql/schemachanger/scbuild/ast_annotator.go b/pkg/sql/schemachanger/scbuild/ast_annotator.go index 1b8d71017e7a..59a2cfad9bf1 100644 --- a/pkg/sql/schemachanger/scbuild/ast_annotator.go +++ b/pkg/sql/schemachanger/scbuild/ast_annotator.go @@ -12,6 +12,9 @@ package scbuild import ( "github.com/cockroachdb/cockroach/pkg/sql/parser" + "github.com/cockroachdb/cockroach/pkg/sql/parser/statements" + "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode" + "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror" "github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scbuild/internal/scbuildstmt" "github.com/cockroachdb/cockroach/pkg/sql/sem/tree" "github.com/cockroachdb/errors" @@ -35,6 +38,13 @@ func newAstAnnotator(original tree.Statement) (*astAnnotator, error) { if err != nil { return nil, err } + // For `CREATE FUNCTION` stmt, overwrite `NumAnnotations` to be the largest of + // parsing the node and parsing all its body statements, to avoid annotation + // index-out-of-bound issue during formatting (https://github.com/cockroachdb/cockroach/issues/104242). + if _, ok := statement.AST.(*tree.CreateFunction); ok { + statement.NumAnnotations = numberAnnotationsOfCreateFunction(statement) + } + return &astAnnotator{ nonExistentNames: map[*tree.TableName]struct{}{}, statement: statement.AST, @@ -88,3 +98,39 @@ func (ann *astAnnotator) ValidateAnnotations() { })) f.FormatNode(ann.statement) } + +// numberAnnotationsOfCreateFunction parse each function body statement +// and return the largest NumAnnotations in `cfStmt` and in each of those +// parsed function body statement. +func numberAnnotationsOfCreateFunction( + cfStmt statements.Statement[tree.Statement], +) tree.AnnotationIdx { + if cf, ok := cfStmt.AST.(*tree.CreateFunction); !ok { + panic(errors.AssertionFailedf("statement is not CREATE FUNCTION")) + } else { + var funcBodyFound bool + var funcBodyStr string + for _, option := range cf.Options { + switch opt := option.(type) { + case tree.FunctionBodyStr: + funcBodyFound = true + funcBodyStr = string(opt) + } + } + if !funcBodyFound { + panic(pgerror.New(pgcode.InvalidFunctionDefinition, "no function body specified")) + } + // ret = max(cfSmt.NumAnnotations, max(funcBodyStmt.NumAnnotations within funcBodyStmts)) + ret := cfStmt.NumAnnotations + funcBodyStmts, err := parser.Parse(funcBodyStr) + if err != nil { + panic(err) + } + for _, funcBodyStmt := range funcBodyStmts { + if funcBodyStmt.NumAnnotations > ret { + ret = funcBodyStmt.NumAnnotations + } + } + return ret + } +}