plpgsql: prevent a crash in an edge case during formatting

This commit fixes an issue that is similar to the one that was fixed in
22dabb0. In particular, each parser
(both SQL one and PLpgSQL one) keeps track of the number of annotations
found in the current stmt. Annotations are used during object name
resolution, and the AST nodes can reference annotation indices. If an
annotation index is set, then it assumes that during formatting the
corresponding annotation will be available in the supplied Annotations
container.

In the PLpgSQL parser we use an ephemeral SQL parser instance to process
SQL statements. Previously, we would simply discard the information about
whether any annotations were found in the SQL stmt which could lead to
undefined behavior down the line. This commit audits all usages of the
SQL parser from PLpgSQL one to ensure that we capture the number of
annotations used in that stmt and create the Annotations container with
that capacity, stored in the AST. This container is then substituted
into the FmtCtx for those ASTs. Note that since we currently don't
support DDLs from PLpgSQL, the annotation slots in the container won't
ever be set, so the contribution of this patch is preventing the crash.

Release note (bug fix): Previously, CockroachDB node could crash when
executing DO stmts when they contain currently unsupported DDL stmts like
CREATE TYPE in non-default configuration (additional logging like the one
controlled via `sql.log.all_statements.enabled` cluster setting needed
to be enabled). This bug was introduced in 25.1 release and is now
fixed.

Author: yuzefovich

pkg/sql/logictest/testdata/logic_test/do

@@ -274,3 +274,16 @@ CALL p();
 NOTICE: Hello, world!
 
 subtest end
+
+subtest regression_143974
+
+statement error pgcode 0A000 unimplemented: CREATE TYPE usage inside a function definition.*\n.*\n.*issue-v/110080
+DO $$
+BEGIN
+  IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'mood') THEN
+    CREATE TYPE mood AS ENUM ('happy', 'sad', 'neutral');
+  END IF;
+END
+$$;
+
+subtest end

pkg/sql/parser/testdata/do

@@ -364,3 +364,45 @@ at or near "EOF": syntax error
 DETAIL: source SQL:
 BEGIN DO
         ^
+
+parse
+DO $$
+BEGIN
+  IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'mood') THEN
+    CREATE TYPE mood AS ENUM ('happy', 'sad', 'neutral');
+  END IF;
+END
+$$;
+----
+DO $$
+BEGIN
+IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'mood') THEN
+	CREATE TYPE mood AS ENUM ('happy', 'sad', 'neutral');
+END IF;
+END;
+$$;
+ -- normalized!
+DO $$
+BEGIN
+IF (NOT (EXISTS (SELECT (1) FROM pg_type WHERE ((typname) = ('mood'))))) THEN
+	CREATE TYPE mood AS ENUM ('happy', 'sad', 'neutral');
+END IF;
+END;
+$$;
+ -- fully parenthesized
+DO $$
+BEGIN
+IF NOT EXISTS (SELECT _ FROM pg_type WHERE typname = '_') THEN
+	CREATE TYPE mood AS ENUM ('happy', 'sad', 'neutral');
+END IF;
+END;
+$$;
+ -- literals removed
+DO $$
+BEGIN
+IF NOT EXISTS (SELECT 1 FROM _ WHERE _ = 'mood') THEN
+	CREATE TYPE _ AS ENUM (_, _, _);
+END IF;
+END;
+$$;
+ -- identifiers removed

pkg/sql/plpgsql/parser/lexer.go

@@ -181,10 +181,12 @@ func (l *lexer) MakeExecSqlStmt() (*plpgsqltree.Execute, error) {
 	if target != nil && sqlStmt.AST.StatementReturnType() != tree.Rows {
 		return nil, pgerror.New(pgcode.Syntax, "INTO used with a command that cannot return data")
 	}
+	ann := tree.MakeAnnotations(sqlStmt.NumAnnotations)
 	return &plpgsqltree.Execute{
-		SqlStmt: sqlStmt.AST,
-		Strict:  haveStrict,
-		Target:  target,
+		SqlStmt:     sqlStmt.AST,
+		Strict:      haveStrict,
+		Target:      target,
+		Annotations: &ann,
 	}, nil
 }
 
@@ -295,10 +297,12 @@ func (l *lexer) MakeFetchOrMoveStmt(isMove bool) (plpgsqltree.Statement, error)
 	}
 	// Move past the semicolon.
 	l.lastPos++
+	ann := tree.MakeAnnotations(sqlStmt.NumAnnotations)
 	return &plpgsqltree.Fetch{
-		Cursor: cursor,
-		Target: target,
-		IsMove: isMove,
+		Cursor:      cursor,
+		Target:      target,
+		IsMove:      isMove,
+		Annotations: &ann,
 	}, nil
 }
 
@@ -396,7 +400,11 @@ func (l *lexer) ParseReturnQuery() (plpgsqltree.Statement, error) {
 	if err != nil {
 		return nil, err
 	}
-	return &plpgsqltree.ReturnQuery{SqlStmt: stmt.AST}, nil
+	ann := tree.MakeAnnotations(stmt.NumAnnotations)
+	return &plpgsqltree.ReturnQuery{
+		SqlStmt:     stmt.AST,
+		Annotations: &ann,
+	}, nil
 }
 
 // peekForExecute checks whether the next token is EXECUTE, used to identify

pkg/sql/plpgsql/parser/plpgsql.y

@@ -6,6 +6,7 @@ import (
 
   "github.com/cockroachdb/cockroach/pkg/build"
   "github.com/cockroachdb/cockroach/pkg/sql/parser"
+  "github.com/cockroachdb/cockroach/pkg/sql/parser/statements"
   "github.com/cockroachdb/cockroach/pkg/sql/scanner"
   "github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
   "github.com/cockroachdb/cockroach/pkg/sql/sem/plpgsqltree"
@@ -173,7 +174,11 @@ func (u *plpgsqlSymUnion) cursorScrollOption() tree.CursorScrollOption {
 }
 
 func (u *plpgsqlSymUnion) sqlStatement() tree.Statement {
-    return u.val.(tree.Statement)
+    return u.val.(statements.Statement[tree.Statement]).AST
+}
+
+func (u *plpgsqlSymUnion) numAnnotations() tree.AnnotationIdx {
+    return u.val.(statements.Statement[tree.Statement]).NumAnnotations
 }
 
 func (u *plpgsqlSymUnion) variables() []plpgsqltree.Variable {
@@ -479,10 +484,12 @@ decl_statement: decl_varname decl_const decl_datatype decl_collate decl_notnull
   }
 | decl_varname opt_scrollable CURSOR decl_cursor_args decl_is_for decl_cursor_query
   {
+    ann := tree.MakeAnnotations($6.numAnnotations())
     $$.val = &plpgsqltree.CursorDeclaration{
       Name: plpgsqltree.Variable($1),
       Scroll: $2.cursorScrollOption(),
       Query: $6.sqlStatement(),
+      Annotations: &ann,
     }
   }
 ;
@@ -510,7 +517,7 @@ decl_cursor_query: stmt_until_semi ';'
     if len(stmts) != 1 {
       return setErr(plpgsqllex, errors.New("expected exactly one SQL statement for cursor"))
     }
-    $$.val = stmts[0].AST
+    $$.val = stmts[0]
   }
 ;
 
@@ -1412,10 +1419,12 @@ stmt_open: OPEN IDENT ';'
     if len(stmts) != 1 {
       return setErr(plpgsqllex, errors.New("expected exactly one SQL statement for cursor"))
     }
+    ann := tree.MakeAnnotations(stmts[0].NumAnnotations)
     $$.val = &plpgsqltree.Open{
       CurVar: plpgsqltree.Variable($2),
       Scroll: $3.cursorScrollOption(),
       Query: stmts[0].AST,
+      Annotations: &ann,
     }
   }
 ;