Merge pull request #153913 from yuzefovich/blathers/backport-release-25.3-151849

release-25.3: plpgsql: fix handling of annotations for DO blocks

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