sql: wrap CASE expressions with parenthesis within PLpgSQL context

Currently, if we generate a CASE expression within the condition of the
IF block in PLpgSQL, it cannot be parsed because it's unclear whether
THEN belongs to CASE or to IF. In order to disambiguate this, we need to
wrap CASE expressions in parenthesis, and we now do so in all
sqlsmith-generated queries whenever we're in PLpgSQL context.

Release note: None

Author: yuzefovich

pkg/internal/sqlsmith/plpgsql.go

@@ -14,6 +14,8 @@ import (
 	"github.com/cockroachdb/errors"
 )
 
+var plpgsqlFlags = tree.FmtParsable | tree.FmtPLpgSQLParen
+
 func (s *Smither) makeRoutineBodyPLpgSQL(
 	params tree.ParamTypes, rTyp *types.T, vol tree.RoutineVolatility,
 ) string {
@@ -25,7 +27,7 @@ func (s *Smither) makeRoutineBodyPLpgSQL(
 	// errors.
 	block := s.makePLpgSQLBlock(scope)
 	block.Body = append(block.Body, s.makePLpgSQLReturn(scope))
-	return "\n" + tree.AsStringWithFlags(s.makePLpgSQLBlock(scope), tree.FmtParsable)
+	return "\n" + tree.AsStringWithFlags(s.makePLpgSQLBlock(scope), plpgsqlFlags, tree.FmtInPLpgSQL(true /* inPLpgSQL */))
 }
 
 func (s *Smither) makePLpgSQLBlock(scope plpgsqlBlockScope) *ast.Block {

pkg/internal/sqlsmith/setup_test.go

@@ -8,8 +8,6 @@ package sqlsmith_test
 import (
 	"context"
 	"flag"
-	"fmt"
-	"strings"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
@@ -92,9 +90,7 @@ func TestGenerateParse(t *testing.T) {
 		t.Fatalf("unknown setting %s", settingName)
 	}
 	settings := setting(rnd)
-	t.Log("setting:", settingName, settings.Options)
 	setupSQL := setup(rnd)
-	t.Log(strings.Join(setupSQL, "\n"))
 	for _, stmt := range setupSQL {
 		db.Exec(t, stmt)
 	}
@@ -119,7 +115,6 @@ func TestGenerateParse(t *testing.T) {
 		if err != nil {
 			t.Fatal(err)
 		}
-		fmt.Print("STMT: ", i, "\n", stmt, ";\n\n")
 		if *flagExec {
 			_, err = conn.ExecContext(ctx, `SET statement_timeout = '9s'`)
 			if err != nil {
@@ -129,7 +124,7 @@ func TestGenerateParse(t *testing.T) {
 				es := err.Error()
 				if !seen[es] {
 					seen[es] = true
-					fmt.Printf("ERR (%d): %v\n", i, err)
+					t.Logf("ERR (%d): %v\n", i, err)
 				}
 			}
 		}

pkg/internal/sqlsmith/sqlsmith.go

@@ -192,6 +192,7 @@ var prettyCfg = func() tree.PrettyCfg {
 	cfg := tree.DefaultPrettyCfg()
 	cfg.LineWidth = 120
 	cfg.Simplify = false
+	cfg.FmtFlags = tree.FmtPLpgSQLParen
 	return cfg
 }()
 
@@ -213,10 +214,10 @@ func (s *Smither) Generate() string {
 		i = 0
 
 		printCfg := prettyCfg
-		fl := tree.FmtParsable
+		fl := plpgsqlFlags
 		if s.postgres {
-			printCfg.FmtFlags = tree.FmtPGCatalog
-			fl = tree.FmtPGCatalog
+			printCfg.FmtFlags = tree.FmtPGCatalog | tree.FmtPLpgSQLParen
+			fl = tree.FmtPGCatalog | tree.FmtPLpgSQLParen
 		}
 		p, err := printCfg.Pretty(stmt)
 		if err != nil {

pkg/sql/sem/plpgsqltree/statements.go

@@ -1121,6 +1121,7 @@ func (s *DoBlock) Format(ctx *tree.FmtCtx) {
 		// Format the body of the DO block separately so that FormatStringDollarQuotes
 		// can examine the resulting string and determine how to quote the block.
 		bodyCtx := ctx.Clone()
+		tree.FmtInPLpgSQL(true /* inPLpgSQL */)(bodyCtx)
 		bodyCtx.FormatNode(s.Block)
 		bodyStr := "\n" + bodyCtx.CloseAndGetString()
 

pkg/sql/sem/tree/expr.go

@@ -1438,6 +1438,13 @@ type CaseExpr struct {
 
 // Format implements the NodeFormatter interface.
 func (node *CaseExpr) Format(ctx *FmtCtx) {
+	if ctx.HasFlags(FmtPLpgSQLParen) && ctx.inPLpgSQL {
+		// In some cases in PLpgSQL context we need to wrap the CASE expression
+		// in parenthesis to make it parsable. We do so only if the caller
+		// requested it.
+		ctx.WriteByte('(')
+		defer ctx.WriteByte(')')
+	}
 	ctx.WriteString("CASE ")
 	if node.Expr != nil {
 		ctx.FormatNode(node.Expr)

pkg/sql/sem/tree/format.go

@@ -201,6 +201,10 @@ const (
 
 	// FmtHideHints skips over any hints.
 	FmtHideHints
+
+	// FmtPLpgSQLParen will wrap some expressions in parenthesis when in PLpgSQL
+	// context. This should only be used in tests.
+	FmtPLpgSQLParen
 )
 
 const genericArityIndicator = "__more__"
@@ -343,6 +347,11 @@ type FmtCtx struct {
 	// indexedTypeFormatter is an optional interceptor for formatting
 	// IDTypeReferences differently than normal.
 	indexedTypeFormatter func(*FmtCtx, *OIDTypeReference)
+
+	// inPLpgSQL, if set, indicates that we're formatting a node within PLpgSQL
+	// context.
+	inPLpgSQL bool
+
 	// small scratch buffer to reduce allocations.
 	scratch [64]byte
 }
@@ -404,6 +413,14 @@ func FmtLocation(loc *time.Location) FmtCtxOption {
 	}
 }
 
+// FmtInPLpgSQL modifies FmtCtx to indicate whether we're in the PLpgSQL
+// context.
+func FmtInPLpgSQL(inPLpgSQL bool) FmtCtxOption {
+	return func(ctx *FmtCtx) {
+		ctx.inPLpgSQL = inPLpgSQL
+	}
+}
+
 // NewFmtCtx creates a FmtCtx; only flags that don't require Annotations
 // can be used.
 func NewFmtCtx(f FmtFlags, opts ...FmtCtxOption) *FmtCtx {
@@ -423,14 +440,15 @@ func NewFmtCtx(f FmtFlags, opts ...FmtCtxOption) *FmtCtx {
 // original.
 func (ctx *FmtCtx) Clone() *FmtCtx {
 	newCtx := fmtCtxPool.Get().(*FmtCtx)
+	newCtx.dataConversionConfig = ctx.dataConversionConfig
+	newCtx.location = ctx.location
 	newCtx.flags = ctx.flags
 	newCtx.ann = ctx.ann
 	newCtx.indexedVarFormat = ctx.indexedVarFormat
 	newCtx.placeholderFormat = ctx.placeholderFormat
 	newCtx.tableNameFormatter = ctx.tableNameFormatter
 	newCtx.indexedTypeFormatter = ctx.indexedTypeFormatter
-	newCtx.dataConversionConfig = ctx.dataConversionConfig
-	newCtx.location = ctx.location
+	newCtx.inPLpgSQL = ctx.inPLpgSQL
 	return newCtx
 }