Commit 2c936219 authored by craig[bot]'s avatar craig[bot]

Merge #34501

34501: tree: improve formatting API r=RaduBerinde a=RaduBerinde

This change improves the FmtCtx API: before, the caller had to save
old flags (or custom formatter callback) and then restore them. The
new API does this automatically:

```
oldFlags := ctx.ChangeFlags(x)          ctx.WithFlags(x, func() {
...                              -->      ...
ctx.ChangeFlags(oldFlags)               })
```

The BenchmarkFormatRandomStatements results are unaffected.

Release note: None
Co-authored-by: default avatarRadu Berinde <[email protected]>
parents df765e3e 7e9287a4
......@@ -167,17 +167,15 @@ func rewriteViewQueryDBNames(table *sqlbase.TableDescriptor, newDB string) error
}
// Re-format to change all DB names to `newDB`.
f := tree.NewFmtCtx(tree.FmtParsable)
var reformatTableNames func(ctx *tree.FmtCtx, tn *tree.TableName)
reformatTableNames = func(ctx *tree.FmtCtx, tn *tree.TableName) {
f.SetReformatTableNames(func(ctx *tree.FmtCtx, tn *tree.TableName) {
// empty catalog e.g. ``"".information_schema.tables` should stay empty.
if tn.CatalogName != "" {
tn.CatalogName = tree.Name(newDB)
}
ctx.WithReformatTableNames(nil)
ctx.FormatNode(tn)
ctx.WithReformatTableNames(reformatTableNames)
}
f.WithReformatTableNames(reformatTableNames)
ctx.WithReformatTableNames(nil, func() {
ctx.FormatNode(tn)
})
})
f.FormatNode(stmt.AST)
table.ViewQuery = f.CloseAndGetString()
return nil
......
......@@ -291,9 +291,7 @@ func scrubStmtStatKey(vt VirtualTabler, key string) (string, bool) {
// Re-format to remove most names.
f := tree.NewFmtCtx(tree.FmtAnonymize)
var reformatFn func(ctx *tree.FmtCtx, tn *tree.TableName)
reformatFn = func(ctx *tree.FmtCtx, tn *tree.TableName) {
reformatFn := func(ctx *tree.FmtCtx, tn *tree.TableName) {
virtual, err := vt.getVirtualTableEntry(tn)
if err != nil || virtual.desc == nil {
ctx.WriteByte('_')
......@@ -304,13 +302,13 @@ func scrubStmtStatKey(vt VirtualTabler, key string) (string, bool) {
newTn := *tn
newTn.CatalogName = "_"
oldFlags := ctx.ChangeFlags(tree.FmtParsable)
ctx.WithReformatTableNames(nil)
ctx.FormatNode(&newTn)
ctx.WithReformatTableNames(reformatFn)
ctx.ChangeFlags(oldFlags)
ctx.WithFlags(tree.FmtParsable, func() {
ctx.WithReformatTableNames(nil, func() {
ctx.FormatNode(&newTn)
})
})
}
f.WithReformatTableNames(reformatFn)
f.SetReformatTableNames(reformatFn)
f.FormatNode(stmt.AST)
return f.CloseAndGetString(), true
}
......
......@@ -75,7 +75,7 @@ func (p *planner) CreateView(ctx context.Context, n *tree.CreateView) (planNode,
// discarded immediately after the traversal because it is not needed further.
{
f := tree.NewFmtCtx(tree.FmtParsable)
f.WithReformatTableNames(
f.SetReformatTableNames(
func(_ *tree.FmtCtx, tn *tree.TableName) {
// Persist the database prefix expansion.
if tn.SchemaName != "" {
......
......@@ -76,7 +76,7 @@ func ConvertToMappedSpecOrdering(
// with their values.
func ExprFmtCtxBase(evalCtx *tree.EvalContext) *tree.FmtCtx {
fmtCtx := tree.NewFmtCtx(tree.FmtCheckEquivalence)
fmtCtx.WithPlaceholderFormat(
fmtCtx.SetPlaceholderFormat(
func(fmtCtx *tree.FmtCtx, p *tree.Placeholder) {
d, err := p.Eval(evalCtx)
if err != nil {
......
......@@ -96,15 +96,13 @@ func MakeExpression(
// We format the expression using the IndexedVar and Placeholder formatting interceptors.
fmtCtx := distsqlpb.ExprFmtCtxBase(evalCtx)
if indexVarMap != nil {
fmtCtx.WithIndexedVarFormat(
func(ctx *tree.FmtCtx, idx int) {
remappedIdx := indexVarMap[idx]
if remappedIdx < 0 {
panic(fmt.Sprintf("unmapped index %d", idx))
}
ctx.Printf("@%d", remappedIdx+1)
},
)
fmtCtx.SetIndexedVarFormat(func(ctx *tree.FmtCtx, idx int) {
remappedIdx := indexVarMap[idx]
if remappedIdx < 0 {
panic(fmt.Sprintf("unmapped index %d", idx))
}
ctx.Printf("@%d", remappedIdx+1)
})
}
fmtCtx.FormatNode(outExpr)
if log.V(1) {
......
......@@ -110,9 +110,9 @@ func (p *planner) makeExplainPlanNodeWithPlan(
// We also avoid calling ctx.FormatNode because when
// types are visible, this would cause the type information
// to be printed twice.
ctx.WithPlaceholderFormat(nil)
placeholder.Format(ctx)
ctx.WithPlaceholderFormat(e.showPlaceholderValues)
ctx.WithPlaceholderFormat(nil, func() {
placeholder.Format(ctx)
})
return
}
ctx.FormatNode(d)
......@@ -389,7 +389,7 @@ func (e *explainer) expr(v observeVerbosity, nodeName, fieldName string, n int,
}
f := tree.NewFmtCtx(e.fmtFlags)
f.WithPlaceholderFormat(e.showPlaceholderValues)
f.SetPlaceholderFormat(e.showPlaceholderValues)
f.FormatNode(expr)
e.attr(nodeName, fieldName, f.CloseAndGetString())
}
......
......@@ -62,7 +62,7 @@ func fmtInterceptor(f *memo.ExprFmtCtx, tp treeprinter.Node, nd opt.Expr) bool {
return false
}
fmtCtx := tree.NewFmtCtx(tree.FmtSimple)
fmtCtx.WithIndexedVarFormat(func(ctx *tree.FmtCtx, idx int) {
fmtCtx.SetIndexedVarFormat(func(ctx *tree.FmtCtx, idx int) {
fullyQualify := !f.HasFlags(memo.ExprFmtHideQualifications)
alias := md.ColumnMeta(opt.ColumnID(idx + 1)).QualifiedAlias(fullyQualify)
ctx.WriteString(alias)
......
......@@ -969,18 +969,18 @@ func (node *Subquery) Format(ctx *FmtCtx) {
} else {
// Ensure that type printing is disabled during the recursion, as
// the type annotations are not available in subqueries.
oldFlags := ctx.ChangeFlags(ctx.flags & ^FmtShowTypes)
if node.Exists {
ctx.WriteString("EXISTS ")
}
if node.Select == nil {
// If the subquery is generated by the optimizer, we
// don't have an actual statement.
ctx.WriteString("<unknown>")
} else {
ctx.FormatNode(node.Select)
}
ctx.ChangeFlags(oldFlags)
ctx.WithFlags(ctx.flags & ^FmtShowTypes, func() {
if node.Exists {
ctx.WriteString("EXISTS ")
}
if node.Select == nil {
// If the subquery is generated by the optimizer, we
// don't have an actual statement.
ctx.WriteString("<unknown>")
} else {
ctx.FormatNode(node.Select)
}
})
}
}
......@@ -1337,9 +1337,9 @@ func (node *FuncExpr) Format(ctx *FmtCtx) {
// We need to remove name anonymization for the function name in
// particular. Do this by overriding the flags.
flagsSave := ctx.ChangeFlags(ctx.flags &^ FmtAnonymize)
ctx.FormatNode(&node.Func)
ctx.ChangeFlags(flagsSave)
ctx.WithFlags(ctx.flags&^FmtAnonymize, func() {
ctx.FormatNode(&node.Func)
})
ctx.WriteByte('(')
ctx.WriteString(typ)
......
......@@ -201,19 +201,31 @@ func NewFmtCtx(f FmtFlags) *FmtCtx {
return ctx
}
// WithReformatTableNames modifies FmtCtx to instructs the pretty-printer
// to substitute the printing of table names using the provided function.
func (ctx *FmtCtx) WithReformatTableNames(fn func(*FmtCtx, *TableName)) *FmtCtx {
ctx.tableNameFormatter = fn
return ctx
// SetReformatTableNames modifies FmtCtx to to substitute the printing of table
// names using the provided function.
func (ctx *FmtCtx) SetReformatTableNames(tableNameFmt func(*FmtCtx, *TableName)) {
ctx.tableNameFormatter = tableNameFmt
}
// ChangeFlags changes the flags in the FmtCtx. The old flags are returned so
// they can be restored.
func (ctx *FmtCtx) ChangeFlags(newFlags FmtFlags) (oldFlags FmtFlags) {
oldFlags = ctx.flags
ctx.flags = newFlags
return oldFlags
// WithReformatTableNames modifies FmtCtx to to substitute the printing of table
// names using the provided function, calls fn, then restores the original table
// formatting.
func (ctx *FmtCtx) WithReformatTableNames(tableNameFmt func(*FmtCtx, *TableName), fn func()) {
old := ctx.tableNameFormatter
ctx.tableNameFormatter = tableNameFmt
defer func() { ctx.tableNameFormatter = old }()
fn()
}
// WithFlags changes the flags in the FmtCtx, runs the given function, then
// restores the old flags.
func (ctx *FmtCtx) WithFlags(flags FmtFlags, fn func()) {
oldFlags := ctx.flags
ctx.flags = flags
defer func() { ctx.flags = oldFlags }()
fn()
}
// HasFlags returns true iff the given flags are set in the formatter context.
......@@ -247,18 +259,25 @@ func FmtExpr(base FmtFlags, showTypes bool, symbolicVars bool, showTableAliases
return base
}
// WithIndexedVarFormat modifies FmtCtx to customize the printing of
// SetIndexedVarFormat modifies FmtCtx to customize the printing of
// IndexedVars using the provided function.
func (ctx *FmtCtx) WithIndexedVarFormat(fn func(ctx *FmtCtx, idx int)) *FmtCtx {
func (ctx *FmtCtx) SetIndexedVarFormat(fn func(ctx *FmtCtx, idx int)) {
ctx.indexedVarFormat = fn
return ctx
}
// WithPlaceholderFormat modifies FmtCtx to customizes the printing of
// SetPlaceholderFormat modifies FmtCtx to customize the printing of
// StarDatums using the provided function.
func (ctx *FmtCtx) WithPlaceholderFormat(placeholderFn func(_ *FmtCtx, _ *Placeholder)) *FmtCtx {
func (ctx *FmtCtx) SetPlaceholderFormat(placeholderFn func(_ *FmtCtx, _ *Placeholder)) {
ctx.placeholderFormat = placeholderFn
return ctx
}
// WithPlaceholderFormat changes the placeholder formatting function, calls the
// given function, then restores the placeholder function.
func (ctx *FmtCtx) WithPlaceholderFormat(placeholderFn func(_ *FmtCtx, _ *Placeholder), fn func()) {
old := ctx.placeholderFormat
ctx.placeholderFormat = placeholderFn
defer func() { ctx.placeholderFormat = old }()
fn()
}
// NodeFormatter is implemented by nodes that can be pretty-printed.
......
......@@ -143,7 +143,7 @@ func TestFormatTableName(t *testing.T) {
}
f := tree.NewFmtCtx(tree.FmtSimple)
f.WithReformatTableNames(func(ctx *tree.FmtCtx, _ *tree.TableName) {
f.SetReformatTableNames(func(ctx *tree.FmtCtx, _ *tree.TableName) {
ctx.WriteString("xoxoxo")
})
......
......@@ -77,11 +77,9 @@ func TestIndexedVars(t *testing.T) {
// Test formatting using the indexed var format interceptor.
f := NewFmtCtx(FmtSimple)
f.WithIndexedVarFormat(
func(ctx *FmtCtx, idx int) {
ctx.Printf("customVar%d", idx)
},
)
f.SetIndexedVarFormat(func(ctx *FmtCtx, idx int) {
ctx.Printf("customVar%d", idx)
})
f.FormatNode(typedExpr)
str = f.CloseAndGetString()
......
......@@ -37,11 +37,11 @@ func (node *SetVar) Format(ctx *FmtCtx) {
ctx.FormatNode(&node.Values)
ctx.WriteString(")")
} else {
oldFlags := ctx.ChangeFlags(ctx.flags & ^FmtAnonymize)
// Session var names never contain PII and should be distinguished
// for feature tracking purposes.
ctx.FormatNameP(&node.Name)
ctx.ChangeFlags(oldFlags)
ctx.WithFlags(ctx.flags & ^FmtAnonymize, func() {
// Session var names never contain PII and should be distinguished
// for feature tracking purposes.
ctx.FormatNameP(&node.Name)
})
ctx.WriteString(" = ")
ctx.FormatNode(&node.Values)
......@@ -59,9 +59,9 @@ func (node *SetClusterSetting) Format(ctx *FmtCtx) {
ctx.WriteString("SET CLUSTER SETTING ")
// Cluster setting names never contain PII and should be distinguished
// for feature tracking purposes.
oldFlags := ctx.ChangeFlags(ctx.flags & ^FmtAnonymize)
ctx.FormatNameP(&node.Name)
ctx.ChangeFlags(oldFlags)
ctx.WithFlags(ctx.flags & ^FmtAnonymize, func() {
ctx.FormatNameP(&node.Name)
})
ctx.WriteString(" = ")
ctx.FormatNode(node.Value)
......
......@@ -35,9 +35,9 @@ func (node *ShowVar) Format(ctx *FmtCtx) {
ctx.WriteString("SHOW ")
// Session var names never contain PII and should be distinguished
// for feature tracking purposes.
oldFlags := ctx.ChangeFlags(ctx.flags & ^FmtAnonymize)
ctx.FormatNameP(&node.Name)
ctx.ChangeFlags(oldFlags)
ctx.WithFlags(ctx.flags & ^FmtAnonymize, func() {
ctx.FormatNameP(&node.Name)
})
}
// ShowClusterSetting represents a SHOW CLUSTER SETTING statement.
......@@ -50,9 +50,9 @@ func (node *ShowClusterSetting) Format(ctx *FmtCtx) {
ctx.WriteString("SHOW CLUSTER SETTING ")
// Cluster setting names never contain PII and should be distinguished
// for feature tracking purposes.
oldFlags := ctx.ChangeFlags(ctx.flags & ^FmtAnonymize)
ctx.FormatNameP(&node.Name)
ctx.ChangeFlags(oldFlags)
ctx.WithFlags(ctx.flags & ^FmtAnonymize, func() {
ctx.FormatNameP(&node.Name)
})
}
// BackupDetails represents the type of details to display for a SHOW BACKUP
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment