util/pretty: mitigate stack overflows of Pretty

This commit reduces the chance of a stack overflow from recursive calls
of `*beExec.be`. The `Pretty` function will now return an internal error
if the recursive depth of `be` surpasses 100,000. While still an
internal error, this is preferable to a stack overflow which will crash
the process.

Informs cockroachdb#91197

Release note: None

pkg/cli/sqlfmt.go

@@ -76,7 +76,11 @@ func runSQLFmt(cmd *cobra.Command, args []string) error {
 	}
 
 	for i := range sl {
-		fmt.Print(cfg.Pretty(sl[i].AST))
+		p, err := cfg.Pretty(sl[i].AST)
+		if err != nil {
+			return err
+		}
+		fmt.Print(p)
 		if len(sl) > 1 {
 			fmt.Print(";")
 		}

pkg/cmd/reduce/reduce/reducesql/reducesql.go

@@ -438,7 +438,11 @@ func joinASTs(stmts []tree.NodeFormatter) string {
 			UseTabs:   false,
 			Simplify:  true,
 		}
-		sb.WriteString(cfg.Pretty(stmt))
+		p, err := cfg.Pretty(stmt)
+		if err != nil {
+			panic(err)
+		}
+		sb.WriteString(p)
 		sb.WriteString(";")
 	}
 	return sb.String()

pkg/internal/sqlsmith/setup_test.go

@@ -123,7 +123,10 @@ func TestGenerateParse(t *testing.T) {
 		if err != nil {
 			t.Fatalf("%v: %v", stmt, err)
 		}
-		stmt = sqlsmith.TestingPrettyCfg.Pretty(parsed.AST)
+		stmt, err = sqlsmith.TestingPrettyCfg.Pretty(parsed.AST)
+		if err != nil {
+			t.Fatal(err)
+		}
 		fmt.Print("STMT: ", i, "\n", stmt, ";\n\n")
 		if *flagExec {
 			db.Exec(t, `SET statement_timeout = '9s'`)

pkg/internal/sqlsmith/sqlsmith.go

@@ -190,7 +190,11 @@ func (s *Smither) Generate() string {
 			continue
 		}
 		i = 0
-		return prettyCfg.Pretty(stmt)
+		p, err := prettyCfg.Pretty(stmt)
+		if err != nil {
+			panic(err)
+		}
+		return p
 	}
 }
 

pkg/sql/explain_bundle.go

@@ -148,7 +148,10 @@ func buildStatementBundle(
 	if plan == nil {
 		return diagnosticsBundle{collectionErr: errors.AssertionFailedf("execution terminated early")}
 	}
-	b := makeStmtBundleBuilder(explainFlags, db, ie, stmtRawSQL, plan, trace, placeholders, sv)
+	b, err := makeStmtBundleBuilder(explainFlags, db, ie, stmtRawSQL, plan, trace, placeholders, sv)
+	if err != nil {
+		return diagnosticsBundle{collectionErr: err, errorStrings: b.errorStrings}
+	}
 
 	b.addStatement()
 	b.addOptPlans(ctx)
@@ -226,18 +229,21 @@ func makeStmtBundleBuilder(
 	trace tracingpb.Recording,
 	placeholders *tree.PlaceholderInfo,
 	sv *settings.Values,
-) stmtBundleBuilder {
+) (stmtBundleBuilder, error) {
 	b := stmtBundleBuilder{
 		flags: flags, db: db, ie: ie, plan: plan, trace: trace, placeholders: placeholders, sv: sv,
 	}
-	b.buildPrettyStatement(stmtRawSQL)
+	err := b.buildPrettyStatement(stmtRawSQL)
+	if err != nil {
+		return stmtBundleBuilder{}, err
+	}
 	b.z.Init()
-	return b
+	return b, nil
 }
 
 // buildPrettyStatement saves the pretty-printed statement (without any
 // placeholder arguments).
-func (b *stmtBundleBuilder) buildPrettyStatement(stmtRawSQL string) {
+func (b *stmtBundleBuilder) buildPrettyStatement(stmtRawSQL string) (err error) {
 	// If we hit an early error, stmt or stmt.AST might not be initialized yet. In
 	// this case use the original raw SQL.
 	if b.plan.stmt == nil || b.plan.stmt.AST == nil {
@@ -255,7 +261,10 @@ func (b *stmtBundleBuilder) buildPrettyStatement(stmtRawSQL string) {
 		cfg.Align = tree.PrettyNoAlign
 		cfg.JSONFmt = true
 		cfg.ValueRedaction = b.flags.RedactValues
-		b.stmt = cfg.Pretty(b.plan.stmt.AST)
+		b.stmt, err = cfg.Pretty(b.plan.stmt.AST)
+		if err != nil {
+			return err
+		}
 
 		// If we had ValueRedaction set, Pretty surrounded all constants with
 		// redaction markers. We must call Redact to fully redact them.
@@ -266,6 +275,7 @@ func (b *stmtBundleBuilder) buildPrettyStatement(stmtRawSQL string) {
 	if b.stmt == "" {
 		b.stmt = "-- no statement"
 	}
+	return nil
 }
 
 // addStatement adds the pretty-printed statement in b.stmt as file

pkg/sql/logictest/logic.go

@@ -728,7 +728,11 @@ func (ls *logicStatement) readSQL(
 					if i > 0 {
 						fmt.Fprintln(&newSyntax, ";")
 					}
-					fmt.Fprint(&newSyntax, pcfg.Pretty(stmtList[i].AST))
+					p, err := pcfg.Pretty(stmtList[i].AST)
+					if err != nil {
+						return "", errors.Wrapf(err, "error while pretty printing")
+					}
+					fmt.Fprint(&newSyntax, p)
 				}
 				return newSyntax.String(), nil
 			}(ls.sql)

pkg/sql/opt/optgen/cmd/optfmt/main.go

@@ -173,7 +173,10 @@ func prettyify(r io.Reader, n int, exprgen bool) (string, error) {
 		exprs = parser.Exprs()
 	}
 	d := p.toDoc(exprs)
-	s := pretty.Pretty(d, n, false, 4, nil)
+	s, err := pretty.Pretty(d, n, false, 4, nil)
+	if err != nil {
+		return "", err
+	}
 
 	// Remove any whitespace at EOL. This can happen in define rules where
 	// we always insert a blank line above comments which are nested with

pkg/sql/sem/builtins/builtins.go

@@ -11157,7 +11157,11 @@ func prettyStatement(p tree.PrettyCfg, stmt string) (string, error) {
 	}
 	var formattedStmt strings.Builder
 	for idx := range stmts {
-		formattedStmt.WriteString(p.Pretty(stmts[idx].AST))
+		p, err := p.Pretty(stmts[idx].AST)
+		if err != nil {
+			return "", err
+		}
+		formattedStmt.WriteString(p)
 		if len(stmts) > 1 {
 			formattedStmt.WriteString(";")
 		}

pkg/sql/sem/tree/pretty.go

@@ -157,13 +157,13 @@ func (p *PrettyCfg) bracketKeyword(
 }
 
 // Pretty pretty prints stmt with default options.
-func Pretty(stmt NodeFormatter) string {
+func Pretty(stmt NodeFormatter) (string, error) {
 	cfg := DefaultPrettyCfg()
 	return cfg.Pretty(stmt)
 }
 
 // Pretty pretty prints stmt with specified options.
-func (p *PrettyCfg) Pretty(stmt NodeFormatter) string {
+func (p *PrettyCfg) Pretty(stmt NodeFormatter) (string, error) {
 	doc := p.Doc(stmt)
 	return pretty.Pretty(doc, p.LineWidth, p.UseTabs, p.TabWidth, p.Case)
 }

pkg/sql/sem/tree/pretty_test.go

@@ -18,6 +18,7 @@ import (
 	"os"
 	"path/filepath"
 	"runtime"
+	"strconv"
 	"strings"
 	"testing"
 
@@ -30,6 +31,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/pretty"
+	"github.com/stretchr/testify/assert"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -112,7 +114,10 @@ func runTestPrettyData(
 				for p := range work {
 					thisCfg := cfg
 					thisCfg.LineWidth = p.numCols
-					res[p.idx] = thisCfg.Pretty(stmt.AST)
+					res[p.idx], err = thisCfg.Pretty(stmt.AST)
+					if err != nil {
+						t.Fatal(err)
+					}
 				}
 				return nil
 			}
@@ -178,14 +183,43 @@ func TestPrettyVerify(t *testing.T) {
 			if err != nil {
 				t.Fatal(err)
 			}
-			got := tree.Pretty(stmt.AST)
+			got, err := tree.Pretty(stmt.AST)
+			if err != nil {
+				t.Fatal(err)
+			}
 			if pretty != got {
 				t.Fatalf("got: %s\nexpected: %s", got, pretty)
 			}
 		})
 	}
 }
 
+func TestPrettyBigStatement(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	// Create a SELECT statement with a 1 million item IN expression. Without
+	// mitigation, this can cause stack overflows - see #91197.
+	var sb strings.Builder
+	sb.WriteString("SELECT * FROM foo WHERE id IN (")
+	for i := 0; i < 1_000_000; i++ {
+		if i != 0 {
+			sb.WriteByte(',')
+		}
+		sb.WriteString(strconv.Itoa(i))
+	}
+	sb.WriteString(");")
+
+	stmt, err := parser.ParseOne(sb.String())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	cfg := tree.DefaultPrettyCfg()
+	_, err = cfg.Pretty(stmt.AST)
+	assert.Errorf(t, err, "max call stack depth of be exceeded")
+}
+
 func BenchmarkPrettyData(b *testing.B) {
 	matches, err := filepath.Glob(datapathutils.TestDataPath(b, "pretty", "*.sql"))
 	if err != nil {
@@ -226,7 +260,10 @@ func TestPrettyExprs(t *testing.T) {
 	}
 
 	for expr, pretty := range tests {
-		got := tree.Pretty(expr)
+		got, err := tree.Pretty(expr)
+		if err != nil {
+			t.Fatal(err)
+		}
 		if pretty != got {
 			t.Fatalf("got: %s\nexpected: %s", got, pretty)
 		}

pkg/sql/show_create_clauses.go

@@ -169,7 +169,12 @@ func formatViewQueryForDisplay(
 				desc.GetName(), desc.GetID(), err)
 			return
 		}
-		query = cfg.Pretty(parsed.AST)
+		query, err = cfg.Pretty(parsed.AST)
+		if err != nil {
+			log.Warningf(ctx, "error printing query for view %s (%v): %+v",
+				desc.GetName(), desc.GetID(), err)
+			return
+		}
 	}()
 
 	typeReplacedViewQuery, err := formatViewQueryTypesForDisplay(ctx, semaCtx, sessionData, desc)

pkg/testutils/sqlutils/pretty.go

@@ -66,7 +66,10 @@ func VerifyStatementPrettyRoundtrip(t *testing.T, sql string) {
 		//
 		origStmt := stmts[i].AST
 		// Be careful to not simplify otherwise the tests won't round trip.
-		prettyStmt := cfg.Pretty(origStmt)
+		prettyStmt, err := cfg.Pretty(origStmt)
+		if err != nil {
+			t.Fatalf("%s: %s", err, prettyStmt)
+		}
 		parsedPretty, err := parser.ParseOne(prettyStmt)
 		if err != nil {
 			t.Fatalf("%s: %s", err, prettyStmt)

pkg/util/pretty/BUILD.bazel

@@ -9,6 +9,7 @@ go_library(
     ],
     importpath = "github.com/cockroachdb/cockroach/pkg/util/pretty",
     visibility = ["//visibility:public"],
+    deps = ["@com_github_cockroachdb_errors//:errors"],
 )
 
 go_test(