Merge pull request #146501 from yuzefovich/blathers/backport-release-25.2-146406

release-25.2: plpgsql: prevent a crash in an edge case during formatting

Author: yuzefovich

pkg/sql/conn_executor_exec.go

@@ -676,6 +676,12 @@ func (ex *connExecutor) execStmtInOpenState(
 	p.semaCtx.Placeholders.Assign(pinfo, stmt.NumPlaceholders)
 	p.extendedEvalCtx.Placeholders = &p.semaCtx.Placeholders
 
+	if buildutil.CrdbTestBuild {
+		// Ensure that each statement is formatted regardless of logging
+		// settings.
+		_ = p.FormatAstAsRedactableString(stmt.AST, p.extendedEvalCtx.Annotations)
+	}
+
 	// This flag informs logging decisions.
 	// Some statements are not dispatched to the execution engine and need
 	// some special plan initialization for logging.
@@ -1673,6 +1679,12 @@ func (ex *connExecutor) execStmtInOpenStateWithPausablePortal(
 	p.semaCtx.Placeholders.Assign(pinfo, vars.stmt.NumPlaceholders)
 	p.extendedEvalCtx.Placeholders = &p.semaCtx.Placeholders
 
+	if buildutil.CrdbTestBuild {
+		// Ensure that each statement is formatted regardless of logging
+		// settings.
+		_ = p.FormatAstAsRedactableString(vars.stmt.AST, p.extendedEvalCtx.Annotations)
+	}
+
 	// This flag informs logging decisions.
 	// Some statements are not dispatched to the execution engine and need
 	// some special plan initialization for logging.

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,
     }
   }
 ;