plpgsql: fix handling of annotations for DO blocks

This commit should fix for good annotations handling for DO blocks that
was recently attempted to be fixed in 3b8ab38.
(That commit is partially reversed too).

The general issue is that in the PLpgSQL parser we use an ephemeral SQL
parser instance to process SQL statements. However, when processing the
DO block stmts we need to preserve the annotations state between each
stmt. This is now achieved by extending the SQL parser infrastructure to
give an option to override the initial annotations index as well as keep
on reusing the counter across multiple stmts. This allows us to
correctly count the number of annotations that we need for the whole DO
block during the initial parsing, and then we use the allocated
annotations container to set the items in the optbuild.

Release note (bug fix): Previously, CockroachDB node could crash when
executing DO stmts when they contain user-defined types (possibly
non-existing) 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/catalog/redact/redact_test.go

@@ -32,15 +32,18 @@ func TestRedactQueries(t *testing.T) {
 	defer srv.Stopper().Stop(ctx)
 	codec := srv.ApplicationLayer().Codec()
 	tdb := sqlutils.MakeSQLRunner(db)
-	tdb.Exec(t, "CREATE TABLE kv (k INT PRIMARY KEY, v STRING)")
-	tdb.Exec(t, "CREATE VIEW view AS SELECT k, v FROM kv WHERE v <> 'constant literal'")
-	tdb.Exec(t, "CREATE TABLE ctas AS SELECT k, v FROM kv WHERE v <> 'constant literal'")
+	tdb.Exec(t, "CREATE TYPE my_enum AS ENUM ('foo', 'bar')")
+	tdb.Exec(t, "CREATE TABLE kv (k INT PRIMARY KEY, v STRING, e my_enum)")
+	tdb.Exec(t, "CREATE VIEW view AS SELECT k, v, e FROM kv WHERE v <> 'constant literal' AND e <> 'foo'")
+	tdb.Exec(t, "CREATE TABLE ctas AS SELECT k, v, e FROM kv WHERE v <> 'constant literal' AND e <> 'foo'")
 	tdb.Exec(t, `
 CREATE FUNCTION f1() RETURNS INT
 LANGUAGE SQL
 AS $$
 SELECT k FROM kv WHERE v != 'foo';
 SELECT k FROM kv WHERE v = 'bar';
+SELECT k FROM kv WHERE e != 'foo';
+SELECT k FROM kv WHERE e = 'bar';
 $$`)
 	tdb.Exec(t, `
 CREATE FUNCTION f2() RETURNS INT
@@ -49,8 +52,9 @@ AS $$
 DECLARE
 x INT := 0;
 y TEXT := 'bar';
+z my_enum;
 BEGIN
-SELECT k FROM kv WHERE v != 'foo';
+SELECT k FROM kv WHERE v != 'foo' AND e != 'bar'::my_enum;
 RETURN x + 3;
 END;
 $$`)
@@ -61,7 +65,7 @@ $$`)
 		)
 		mut := tabledesc.NewBuilder(view.TableDesc()).BuildCreatedMutableTable()
 		require.Empty(t, redact.Redact(mut.DescriptorProto()))
-		require.Equal(t, `SELECT k, v FROM defaultdb.public.kv WHERE v != '_'`, mut.ViewQuery)
+		require.Equal(t, `SELECT k, v, e FROM defaultdb.public.kv WHERE (v != '_') AND (e != '_')`, mut.ViewQuery)
 	})
 
 	t.Run("create table as", func(t *testing.T) {
@@ -70,14 +74,14 @@ $$`)
 		)
 		mut := tabledesc.NewBuilder(ctas.TableDesc()).BuildCreatedMutableTable()
 		require.Empty(t, redact.Redact(mut.DescriptorProto()))
-		require.Equal(t, `SELECT k, v FROM defaultdb.public.kv WHERE v != '_'`, mut.CreateQuery)
+		require.Equal(t, `SELECT k, v, e FROM defaultdb.public.kv WHERE (v != '_') AND (e != '_')`, mut.CreateQuery)
 	})
 
 	t.Run("create function sql", func(t *testing.T) {
 		fn := desctestutils.TestingGetFunctionDescriptor(kvDB, codec, "defaultdb", "public", "f1")
 		mut := funcdesc.NewBuilder(fn.FuncDesc()).BuildCreatedMutableFunction()
 		require.Empty(t, redact.Redact(mut.DescriptorProto()))
-		require.Equal(t, `SELECT k FROM defaultdb.public.kv WHERE v != '_'; SELECT k FROM defaultdb.public.kv WHERE v = '_';`, mut.FunctionBody)
+		require.Equal(t, `SELECT k FROM defaultdb.public.kv WHERE v != '_'; SELECT k FROM defaultdb.public.kv WHERE v = '_'; SELECT k FROM defaultdb.public.kv WHERE e != '_'; SELECT k FROM defaultdb.public.kv WHERE e = '_';`, mut.FunctionBody)
 	})
 
 	t.Run("create function plpgsql", func(t *testing.T) {
@@ -87,8 +91,9 @@ $$`)
 		require.Equal(t, `DECLARE
 x INT8 := _;
 y STRING := '_';
+z @100104;
 BEGIN
-SELECT k FROM defaultdb.public.kv WHERE v != '_';
+SELECT k FROM defaultdb.public.kv WHERE (v != '_') AND (e != '_':::@100104);
 RETURN x + _;
 END;
 `, mut.FunctionBody)

pkg/sql/logictest/testdata/logic_test/do

@@ -285,3 +285,23 @@ END
 $$;
 
 subtest end
+
+subtest regression_151848
+
+statement error pgcode 42704 type \"bar\" does not exist
+DO $$
+DECLARE
+  foo bar;
+BEGIN
+END;
+$$;
+
+statement error pgcode 42704 type \"_\" does not exist
+DO $$
+DECLARE
+  foo _[];
+BEGIN
+END;
+$$;
+
+subtest end

pkg/sql/opt/optbuilder/routine.go

@@ -840,8 +840,19 @@ func (b *Builder) buildDo(do *tree.DoBlock, inScope *scope) *scope {
 	// TODO(drewk): Enable memo reuse with DO statements.
 	b.DisableMemoReuse = true
 
-	defer func(oldInsideFuncDep bool) { b.insideFuncDef = oldInsideFuncDep }(b.insideFuncDef)
+	doBlockImpl, ok := do.Code.(*plpgsqltree.DoBlock)
+	if !ok {
+		panic(errors.AssertionFailedf("expected a plpgsql block"))
+	}
+
+	defer func(oldInsideFuncDep bool, oldAnn tree.Annotations) {
+		b.insideFuncDef = oldInsideFuncDep
+		b.semaCtx.Annotations = oldAnn
+		b.evalCtx.Annotations = &b.semaCtx.Annotations
+	}(b.insideFuncDef, b.semaCtx.Annotations)
 	b.insideFuncDef = true
+	b.semaCtx.Annotations = doBlockImpl.Annotations
+	b.evalCtx.Annotations = &b.semaCtx.Annotations
 
 	// Build the routine body.
 	var bodyStmts []string
@@ -850,10 +861,6 @@ func (b *Builder) buildDo(do *tree.DoBlock, inScope *scope) *scope {
 		fmtCtx.FormatNode(do.Code)
 		bodyStmts = []string{fmtCtx.CloseAndGetString()}
 	}
-	doBlockImpl, ok := do.Code.(*plpgsqltree.DoBlock)
-	if !ok {
-		panic(errors.AssertionFailedf("expected a plpgsql block"))
-	}
 	bodyScope := b.buildPLpgSQLDoBody(doBlockImpl)
 
 	// Build a CALL expression that invokes the routine.

pkg/sql/parser/BUILD.bazel

@@ -37,6 +37,7 @@ go_library(
         "//pkg/sql/sem/tree/treecmp",  # keep
         "//pkg/sql/sem/tree/treewindow",  # keep
         "//pkg/sql/types",
+        "//pkg/util/buildutil",
         "//pkg/util/errorutil/unimplemented",
         "//pkg/util/vector",  # keep
         "@com_github_cockroachdb_errors//:errors",

pkg/sql/parser/lexer.go

@@ -38,13 +38,17 @@ type lexer struct {
 	lastError error
 }
 
-func (l *lexer) init(sql string, tokens []sqlSymType, nakedIntType *types.T) {
+// numAnnotations indicates the number of annotations that have already been
+// claimed.
+func (l *lexer) init(
+	sql string, tokens []sqlSymType, nakedIntType *types.T, numAnnotations tree.AnnotationIdx,
+) {
 	l.in = sql
 	l.tokens = tokens
 	l.lastPos = -1
 	l.stmt = nil
 	l.numPlaceholders = 0
-	l.numAnnotations = 0
+	l.numAnnotations = numAnnotations
 	l.lastError = nil
 
 	l.nakedIntType = nakedIntType

pkg/sql/parser/lexer_test.go

@@ -36,7 +36,7 @@ func TestLexer(t *testing.T) {
 			scanTokens = append(scanTokens, lval)
 		}
 		var l lexer
-		l.init(d.sql, scanTokens, defaultNakedIntType)
+		l.init(d.sql, scanTokens, defaultNakedIntType, 0 /* numAnnotations */)
 		var lexTokens []int
 		for {
 			var lval sqlSymType

pkg/sql/parser/parse.go

@@ -25,6 +25,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/scanner"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
+	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/errors"
 )
 
@@ -47,6 +48,7 @@ type Parser struct {
 type ParseOptions struct {
 	intType        *types.T
 	retainComments bool
+	numAnnotations *tree.AnnotationIdx
 }
 
 var DefaultParseOptions = ParseOptions{
@@ -64,6 +66,18 @@ func (po ParseOptions) WithIntType(t *types.T) ParseOptions {
 	return po
 }
 
+// WithNumAnnotations overrides how annotations are handled. If this option is
+// used, then
+// - the given integer indicates the number of annotations already claimed, and
+// - if multiple stmts are parsed, then unique annotation indexes are used
+// across all stmts.
+//
+// When this option is not used, then each stmt is parsed indepedently.
+func (po ParseOptions) WithNumAnnotations(numAnnotations tree.AnnotationIdx) ParseOptions {
+	po.numAnnotations = &numAnnotations
+	return po
+}
+
 // INT8 is the historical interpretation of INT. This should be left
 // alone in the future, since there are many sql fragments stored
 // in various descriptors. Any user input that was created after
@@ -168,14 +182,29 @@ func (p *Parser) parseWithDepth(
 		p.scanner.RetainComments()
 	}
 	defer p.scanner.Cleanup()
+	var numAnnotations tree.AnnotationIdx
+	if options.numAnnotations != nil {
+		numAnnotations = *options.numAnnotations
+	}
 	for {
 		sql, tokens, done := p.scanOneStmt()
-		stmt, err := p.parse(depth+1, sql, tokens, options.intType)
+		stmt, err := p.parse(depth+1, sql, tokens, options.intType, numAnnotations)
 		if err != nil {
 			return nil, err
 		}
 		if stmt.AST != nil {
 			stmts = append(stmts, stmt)
+			if options.numAnnotations != nil {
+				if buildutil.CrdbTestBuild && numAnnotations > stmt.NumAnnotations {
+					return nil, errors.AssertionFailedf(
+						"annotation index has regressed: numAnnotations=%d, stmt.NumAnnotations=%d ",
+						numAnnotations, stmt.NumAnnotations,
+					)
+				}
+				// If this stmt used any annotations, we need to advance the
+				// number of annotations accordingly.
+				numAnnotations = stmt.NumAnnotations
+			}
 		}
 		if done {
 			break
@@ -186,9 +215,13 @@ func (p *Parser) parseWithDepth(
 
 // parse parses a statement from the given scanned tokens.
 func (p *Parser) parse(
-	depth int, sql string, tokens []sqlSymType, nakedIntType *types.T,
+	depth int,
+	sql string,
+	tokens []sqlSymType,
+	nakedIntType *types.T,
+	numAnnotations tree.AnnotationIdx,
 ) (statements.Statement[tree.Statement], error) {
-	p.lexer.init(sql, tokens, nakedIntType)
+	p.lexer.init(sql, tokens, nakedIntType, numAnnotations)
 	defer p.lexer.cleanup()
 	if p.parserImpl.Parse(&p.lexer) != 0 {
 		if p.lexer.lastError == nil {
@@ -371,27 +404,32 @@ func ParseTablePattern(sql string) (tree.TablePattern, error) {
 // the results are undefined if the string contains invalid SQL
 // syntax.
 func ParseExprs(exprs []string) (tree.Exprs, error) {
-	return ParseExprsWithOptions(exprs, DefaultParseOptions)
+	res, _, err := ParseExprsWithOptions(exprs, DefaultParseOptions)
+	return res, err
 }
 
 // ParseExprsWithOptions parses a comma-delimited sequence of SQL scalar
 // expressions with the provided options. The caller is responsible for
 // ensuring that the input is, in fact, a comma-delimited sequence of SQL
 // scalar expressions — the results are undefined if the string contains
 // invalid SQL syntax.
-func ParseExprsWithOptions(exprs []string, opts ParseOptions) (tree.Exprs, error) {
+//
+// It also returns the number of annotations used.
+func ParseExprsWithOptions(
+	exprs []string, opts ParseOptions,
+) (tree.Exprs, tree.AnnotationIdx, error) {
 	if len(exprs) == 0 {
-		return tree.Exprs{}, nil
+		return tree.Exprs{}, 0, nil
 	}
 	stmt, err := ParseOneWithOptions(fmt.Sprintf("SET ROW (%s)", strings.Join(exprs, ",")), opts)
 	if err != nil {
-		return nil, err
+		return nil, 0, err
 	}
 	set, ok := stmt.AST.(*tree.SetVar)
 	if !ok {
-		return nil, errors.AssertionFailedf("expected a SET statement, but found %T", stmt)
+		return nil, 0, errors.AssertionFailedf("expected a SET statement, but found %T", stmt)
 	}
-	return set.Values, nil
+	return set.Values, stmt.NumAnnotations, nil
 }
 
 // ParseExpr parses a SQL scalar expression. The caller is responsible

pkg/sql/parser/testdata/do

@@ -406,3 +406,45 @@ END IF;
 END;
 $$;
  -- identifiers removed
+
+parse
+DO $$
+DECLARE
+  foo bar;
+  _ _[];
+BEGIN
+END;
+$$;
+----
+DO $$
+DECLARE
+foo bar;
+_ _[];
+BEGIN
+END;
+$$;
+ -- normalized!
+DO $$
+DECLARE
+foo bar;
+_ _[];
+BEGIN
+END;
+$$;
+ -- fully parenthesized
+DO $$
+DECLARE
+foo bar;
+_ _[];
+BEGIN
+END;
+$$;
+ -- literals removed
+DO $$
+DECLARE
+_ _;
+_ _[];
+BEGIN
+END;
+$$;
+ -- identifiers removed